cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-09-30 13:52:10 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'polarization' of github.mit.edu:Scattering/CDTools into polarization

Browse Source
This commit is contained in:
Abe Levitan
2021-08-25 17:06:34 -04:00
parent ebafc44acb e7c15273fa
commit cd6a811a63
6 changed files with 522 additions and 281 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/models/base.py
+70 -58
View File
@@ -11,7 +11,7 @@ simulate_to_dataset
Creates a CDataset from the simulation defined in the model
save_results
Saves out a dictionary with the recovered parameters
Simulation
----------
@@ -56,7 +56,7 @@ class CDIModel(t.nn.Module):
def __init__(self):
super(CDIModel,self).__init__()
self.iteration_count = 0
def from_dataset(self, dataset):
raise NotImplementedError()
@@ -79,11 +79,11 @@ class CDIModel(t.nn.Module):
def forward(self, *args):
"""The complete forward model
This model relies on composing the interaction, forward propagator,
and measurement functions which are required to be defined by all
subclasses. It therefore should not be redefined by the subclasses.
The arguments to this function, for any given subclass, will be
the same as the arguments to the interaction function.
"""
@@ -99,7 +99,7 @@ class CDIModel(t.nn.Module):
def simulate_to_dataset(self, args_list):
raise NotImplementedError()
def save_results(self):
raise NotImplementedError()
@@ -107,10 +107,10 @@ class CDIModel(t.nn.Module):
scheduler=None, regularization_factor=None, thread=True,
calculation_width=10):
"""Runs a round of reconstruction using the provided optimizer
This is the basic automatic differentiation reconstruction tool
which all the other, algorithm-specific tools, use.
Like all the other optimization routines, it is defined as a
generator function which yields the average loss each epoch.
@@ -135,7 +135,7 @@ class CDIModel(t.nn.Module):
normalization = 0
for inputs, patterns in data_loader:
normalization += t.sum(patterns).cpu().numpy()
def run_iteration(stop_event=None):
loss = 0
N = 0
@@ -144,7 +144,7 @@ class CDIModel(t.nn.Module):
N += 1
def closure():
optimizer.zero_grad()
input_chunks = [[inp[i:i + calculation_width]
for inp in inputs]
for i in range(0, len(inputs[0]),
@@ -152,7 +152,7 @@ class CDIModel(t.nn.Module):
pattern_chunks = [patterns[i:i + calculation_width]
for i in range(0, len(inputs[0]),
calculation_width)]
total_loss = 0
for inp, pats in zip(input_chunks, pattern_chunks):
# This is just used to allow graceful exit when
@@ -168,7 +168,7 @@ class CDIModel(t.nn.Module):
loss = self.loss(pats,sim_patterns)
loss.backward()
total_loss += loss.detach()
if regularization_factor is not None \
@@ -178,7 +178,7 @@ class CDIModel(t.nn.Module):
return total_loss
loss += optimizer.step(closure).detach().cpu().numpy()
loss /= normalization
if scheduler is not None:
scheduler.step(loss)
@@ -198,7 +198,7 @@ class CDIModel(t.nn.Module):
# If something bad happens, put the exception into the
# result queue
result_queue.put(e)
for it in range(iterations):
if thread:
calc = threading.Thread(target=target, name='calculator', daemon=True)
@@ -206,10 +206,10 @@ class CDIModel(t.nn.Module):
calc.start()
while calc.is_alive():
if hasattr(self, 'figs'):
self.figs[0].canvas.start_event_loop(0.01)
self.figs[0].canvas.start_event_loop(0.01)
else:
calc.join()
except KeyboardInterrupt as e:
stop_event.set()
print('\nAsking execution thread to stop cleanly - please be patient.')
@@ -233,7 +233,7 @@ class CDIModel(t.nn.Module):
regularization_factor=None, thread=True,
calculation_width=10):
"""Runs a round of reconstruction using the Adam optimizer
This is generally accepted to be the most robust algorithm for use
with ptychography. Like all the other optimization routines,
it is defined as a generator function, which yields the average
@@ -267,7 +267,7 @@ class CDIModel(t.nn.Module):
if type(subset) == type(1):
subset = [subset]
dataset = torchdata.Subset(dataset, subset)
# Make a dataloader
data_loader = torchdata.DataLoader(dataset, batch_size=batch_size,
shuffle=True)
@@ -290,17 +290,17 @@ class CDIModel(t.nn.Module):
calculation_width=calculation_width)
def LBFGS_optimize(self, iterations, dataset,
def LBFGS_optimize(self, iterations, dataset,
lr=0.1,history_size=2, subset=None,
regularization_factor=None, thread=True,
calculation_width=10):
"""Runs a round of reconstruction using the L-BFGS optimizer
This algorithm is often less stable that Adam, however in certain
situations or geometries it can be shockingly efficient. Like all
the other optimization routines, it is defined as a generator
function which yields the average loss each epoch.
Note: There is no batch size, because it is a usually a bad idea to use
LBFGS on anything but all the data at onece
@@ -320,14 +320,14 @@ class CDIModel(t.nn.Module):
Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
thread : bool
Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
"""
if subset is not None:
# if just one pattern, turn into a list for convenience
if type(subset) == type(1):
subset = [subset]
dataset = torchdata.Subset(dataset, subset)
# Make a dataloader. This basically does nothing but load all the
# data at once
data_loader = torchdata.DataLoader(dataset, batch_size=len(dataset))
@@ -338,7 +338,7 @@ class CDIModel(t.nn.Module):
lr = lr, history_size=history_size)
#optimizer = MyLBFGS(self.parameters(),
# lr = lr, history_size=history_size)
return self.AD_optimize(iterations, data_loader, optimizer,
regularization_factor=regularization_factor,
thread=thread,
@@ -350,7 +350,7 @@ class CDIModel(t.nn.Module):
nesterov=False, subset=None, regularization_factor=None,
thread=True, calculation_width=10):
"""Runs a round of reconstruction using the SGDoptimizer
This algorithm is often less stable that Adam, but it is simpler
and is the basic workhorse of gradience descent.
@@ -382,7 +382,7 @@ class CDIModel(t.nn.Module):
if type(subset) == type(1):
subset = [subset]
dataset = torchdata.Subset(dataset, subset)
# Make a dataloader
if batch_size is not None:
data_loader = torchdata.DataLoader(dataset, batch_size=batch_size,
@@ -418,28 +418,28 @@ class CDIModel(t.nn.Module):
self.latest_iteration_time + str(self.latest_loss)
else:
return 'No reconstruction iterations performed yet!'
# By default, the plot_list is empty
plot_list = []
def inspect(self, dataset=None, update=True):
"""Plots all the plots defined in the model's plot_list attribute
If update is set to True, it will update any previously plotted set
of plots, if one exists, and then redraw them. Otherwise, it will
plot a new set, and any subsequent updates will update the new set
Optionally, a dataset can be passed, which then will plot any
registered plots which need to incorporate some information from
the dataset (such as geometry or a comparison with measured data).
Plots can be registered in any subclass by defining the plot_list
attribute. This should be a list of tuples in the following format:
( 'Plot Title', function_to_generate_plot(self),
( 'Plot Title', function_to_generate_plot(self),
function_to_determine_whether_to_plot(self))
Where the third element in the tuple (a function that returns
Where the third element in the tuple (a function that returns
True if the plot is relevant) is not required.
Parameters
@@ -448,9 +448,19 @@ class CDIModel(t.nn.Module):
Optional, a dataset matched to the model type
update : bool
Default True, whether to update existing plots or plot new ones
"""
print('base models inspect: checking the object')
a = self.obj.detach()
def saveobj(a, filename):
a = np.abs(a)
plt.imshow(a)
plt.savefig(filename)
f = ['base_a.png', 'base_b.png', 'base_c.png', 'base_d.png']
comp = [a[i, j, :, :] for i, j in zip([0, 0, 1, 1], [0, 1, 0, 1])]
for i in range(4):
saveobj(comp[i], f[i])
first_update = False
if update and hasattr(self, 'figs') and self.figs:
figs = self.figs
@@ -476,7 +486,7 @@ class CDIModel(t.nn.Module):
plotter = plots[1]
if figs is None:
fig = plt.figure()
fig = plt.figure()
self.figs.append(fig)
else:
fig = figs[idx]
@@ -484,13 +494,13 @@ class CDIModel(t.nn.Module):
try:
plotter(self,fig)
plt.title(name)
except TypeError as e:
if dataset is not None:
try:
plotter(self, fig, dataset)
plt.title(name)
except (IndexError, KeyError, AttributeError, np.linalg.LinAlgError) as e:
pass
@@ -498,15 +508,15 @@ class CDIModel(t.nn.Module):
pass
idx += 1
if update:
plt.draw()
fig.canvas.start_event_loop(0.001)
if first_update:
plt.pause(0.05 * len(self.figs))
def save_figures(self, prefix='', extension='.eps'):
"""Saves all currently open inspection figures.
@@ -520,7 +530,7 @@ class CDIModel(t.nn.Module):
By default, the files will be named by the figure titles as defined
in the plot_list. Files can be saved with any extension suported by
matplotlib.pyplot.savefig.
Parameters
----------
prefix : str
@@ -528,7 +538,7 @@ class CDIModel(t.nn.Module):
extention : strategy
Default is .eps, the file extension to save with.
"""
if hasattr(self, 'figs') and self.figs:
figs = self.figs
else:
@@ -538,21 +548,21 @@ class CDIModel(t.nn.Module):
fig.savefig(prefix + fig.axes[0].get_title() + extension,
bbox_inches = 'tight')
def compare(self, dataset):
def compare(self, dataset, logarithmic=False):
"""Opens a tool for comparing simulated and measured diffraction patterns
Parameters
----------
dataset : CDataset
A dataset containing the simulated diffraction patterns to compare against
"""
fig, axes = plt.subplots(1,3,figsize=(12,5.3))
fig.tight_layout(rect=[0.02, 0.09, 0.98, 0.96])
axslider = plt.axes([0.15,0.06,0.75,0.03])
def update_colorbar(im):
# If the update brought the colorbar out of whack
# (say, from clicking back in the navbar)
@@ -564,7 +574,7 @@ class CDIModel(t.nn.Module):
if hasattr(im, 'norecurse') and im.norecurse:
im.norecurse=False
return
im.norecurse=True
im.set_clim(vmin=np.min(im.get_array()),vmax=np.max(im.get_array()))
@@ -572,7 +582,7 @@ class CDIModel(t.nn.Module):
idx = int(idx) % len(dataset)
fig.pattern_idx = idx
updating = True if len(axes[0].images) >= 1 else False
inputs, output = dataset[idx]
sim_data = self.forward(*inputs).detach().cpu().numpy()
sim_data = sim_data
@@ -581,7 +591,11 @@ class CDIModel(t.nn.Module):
mask = self.mask.detach().cpu().numpy()
else:
mask = 1
if logarithmic:
sim_data =np.log(sim_data)/np.log(10)
meas_data = np.log(meas_data)/np.log(10)
if not updating:
axes[0].set_title('Simulated')
axes[1].set_title('Measured')
@@ -600,7 +614,7 @@ class CDIModel(t.nn.Module):
cb3 = plt.colorbar(diff, ax=axes[2], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.1,fraction=0.1)
cb3.ax.tick_params(labelrotation=20)
cb3.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(diff))
else:
sim = axes[0].images[-1]
@@ -614,8 +628,8 @@ class CDIModel(t.nn.Module):
diff = axes[2].images[-1]
diff.set_data((sim_data-meas_data) * mask)
update_colorbar(diff)
# This is dumb but the slider doesn't work unless a reference to it is
# kept somewhere...
self.slider = Slider(axslider, 'Pattern #', 0, len(dataset)-1, valstep=1, valfmt="%d")
@@ -626,7 +640,7 @@ class CDIModel(t.nn.Module):
event.button = None
if not hasattr(event, 'key'):
event.key = None
if event.key == 'up' or event.button == 'up':
update(fig.pattern_idx - 1)
elif event.key == 'down' or event.button == 'down':
@@ -637,5 +651,3 @@ class CDIModel(t.nn.Module):
fig.canvas.mpl_connect('key_press_event',on_action)
fig.canvas.mpl_connect('scroll_event',on_action)
update(0)
CDTools/models/polarized_fancy_ptycho.py
+218 -43
View File
@@ -21,13 +21,13 @@ class PolarizedFancyPtycho(FancyPtycho):
detector_slice=None,
surface_normal=np.array([0.,0.,1.]),
min_translation = t.Tensor([0,0]),
background = None, translation_offsets=None,
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, obj_support=None, oversampling=1,
loss='amplitude mse',units='um'):
super(FancyPtycho, self).__init__(wavelength, detector_geometry,
probe_basis,
probe_guess, obj_guess,
@@ -48,30 +48,205 @@ class PolarizedFancyPtycho(FancyPtycho):
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, opt_for_fft=False, loss='amplitude mse', units='um', left_polarized=True):
model = FancyPtycho.from_dataset(dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um')
model = FancyPtycho.from_dataset(dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um', left_polarized=True)
# Mutate the class to its subclass
# Mutate the class to its subclass
model.__class__ = cls
if left_polarized:
x = 1j
else:
x = -1j
model.probe.data = t.stack((model.probe.data.to(dtype=t.cfloat), x * model.probe.data.to(dtype=t.cfloat)), dim=-3)
obj = t.stack((model.obj.data, model.obj.data), dim=-3)
model.obj.data = t.stack((obj, obj), dim=-4)
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
# 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()
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.show()
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)]
# WHAT IS INDEX?
@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, opt_for_fft=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,
opt_for_fft=opt_for_fft,
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
@@ -84,9 +259,9 @@ class PolarizedFancyPtycho(FancyPtycho):
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
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
@@ -106,7 +281,7 @@ class PolarizedFancyPtycho(FancyPtycho):
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, :, :, :]
@@ -122,7 +297,7 @@ class PolarizedFancyPtycho(FancyPtycho):
def backward_propagator(self, wavefields):
return tools.propagators.inverse_far_field(wavefields)
def measurement(self, wavefields):
wavefields_x = wavefields[..., 0, :, :]
wavefields_y = wavefields[..., 1, :, :]
@@ -145,15 +320,15 @@ class PolarizedFancyPtycho(FancyPtycho):
# 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',
@@ -166,16 +341,16 @@ class PolarizedFancyPtycho(FancyPtycho):
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)
@@ -187,14 +362,14 @@ class PolarizedFancyPtycho(FancyPtycho):
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:
@@ -205,18 +380,18 @@ class PolarizedFancyPtycho(FancyPtycho):
else:
return np.array([np.eye(self.probe.shape[0])]*self.weights.shape[0],
dtype=np.complex64)
def tidy_probes(self, normalization=1, normalize=False):
"""Tidies up the probes
What we want to do here is use all the information on all the probes
to calculate a natural basis for the experiment, and update all the
density matrices to operate in that updated basis
"""
# First we treat the purely incoherent case
# I don't love this pattern of using an if statement with a return
# to catch this case, but because it's so much simpler than the
# unified mode case I think it's appropriate
@@ -228,11 +403,11 @@ class PolarizedFancyPtycho(FancyPtycho):
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()
@@ -248,7 +423,7 @@ class PolarizedFancyPtycho(FancyPtycho):
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
@@ -263,14 +438,14 @@ class PolarizedFancyPtycho(FancyPtycho):
# 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)
@@ -287,32 +462,32 @@ class PolarizedFancyPtycho(FancyPtycho):
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'),
@@ -333,7 +508,7 @@ class PolarizedFancyPtycho(FancyPtycho):
('% 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',
('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)),
@@ -343,7 +518,7 @@ class PolarizedFancyPtycho(FancyPtycho):
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()
@@ -352,7 +527,7 @@ class PolarizedFancyPtycho(FancyPtycho):
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,