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

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,
CDTools/tools/interactions/interactions.py
+52 -52
View File
@@ -17,14 +17,14 @@ __all__ = ['translations_to_pixel', 'pixel_to_translations',
def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.])):
"""Takes real space translations and outputs them in pixel space
This works for any 2D ptychography geometry. It takes in
A set of translations in (x,y) space and outputs the same translations
in internal pixel units perpendicular to the detector.
in internal pixel units perpendicular to the detector.
It uses information on the wavefield basis and, if defined, the
sample normal, to perform the conversion.
The assumed geometry is incoming radiation with a wavevector parallel
to the +z axis, [0,0,1]. The default sample orientation has a surface
normal parallel to this direction
@@ -33,7 +33,7 @@ def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.
----------
basis : torch.Tensor
The real space basis the wavefields are defined in
translations : torch.Tensor
translations : torch.Tensor
A Jx3 stack of real-space translations, or a single translation
surface_normal : torch.Tensor
Optional, the sample's surface normal
@@ -68,18 +68,18 @@ def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.
return pixel_translations[0]
else:
return pixel_translations
def pixel_to_translations(basis, pixel_translations, surface_normal=t.Tensor([0,0,1])):
"""Takes pixel-space translations and outputs them in real space
This works for any 2D ptychography geometry. It takes in
A set of internal pixel unit translations in (i,j) space and
outputs the same translations real (x,y) space
It uses information on the wavefield basis and, if defined, the
sample normal, to perform the conversion.
The assumed geometry is incoming radiation with a wavevector parallel
to the +z axis, [0,0,1]. The default sample orientation has a surface
normal parallel to this direction. Because of this, the z direction
@@ -96,7 +96,7 @@ def pixel_to_translations(basis, pixel_translations, surface_normal=t.Tensor([0,
Returns
-------
real_translations : torch.Tensor
real_translations : torch.Tensor
A Jx3 stack of real-space translations, or a single translation
"""
projection_1 = t.Tensor([[1,0,0],
@@ -129,7 +129,7 @@ def pixel_to_translations(basis, pixel_translations, surface_normal=t.Tensor([0,
def project_translations_to_sample(sample_basis, translations):
"""Takes real space translations and outputs them in pixels in a sample basis
This projection function is designed for the Bragg2DPtycho class. More
broadly, it works to take a set of translations in the lab frame and
convert each one into two values. First, an (i,j) value in pixels
@@ -145,7 +145,7 @@ def project_translations_to_sample(sample_basis, translations):
relative amount the probe needs to be propagated to reach any given
location), a positive motion along the z-axis of the probe forming optics
will lead to a negative propagation distance.
The assumed geometry is incoming radiation with a wavevector parallel
to the +z axis, [0,0,1].
@@ -153,7 +153,7 @@ def project_translations_to_sample(sample_basis, translations):
----------
sample_basis : torch.Tensor
The real space basis the wavefields are defined in
translations : torch.Tensor
translations : torch.Tensor
A Jx3 stack of real-space translations, or a single translation
Returns
@@ -171,7 +171,7 @@ def project_translations_to_sample(sample_basis, translations):
# Then we calculate a matrix which can do the projection
propagation_dir = t.Tensor(np.array([0,0,1])).to(
device=surface_normal.device,
dtype=surface_normal.dtype)
@@ -179,13 +179,13 @@ def project_translations_to_sample(sample_basis, translations):
I = t.eye(3).to(
device=surface_normal.device,
dtype=surface_normal.dtype)
# Here we're setting up a matrix-vector equation mat*answer=input
# At some point ger will need to be replaced by outer, but for now
# outer many places still don't have new enough versions of torch.
mat = t.cat((I - t.ger(propagation_dir,propagation_dir),
surface_normal.unsqueeze(0)))
# And we invert the matrix to do the projection
projector = t.pinverse(mat)[:,:3].to(device=translations.device,
dtype=translations.dtype)
@@ -200,7 +200,7 @@ def project_translations_to_sample(sample_basis, translations):
device=translations.device,
dtype=translations.dtype)
sample_projection = t.mm(basis_vectors_inv, projector).t()
prop_projection = t.mm(propagation_dir_inv, projector).t()
@@ -218,9 +218,9 @@ def project_translations_to_sample(sample_basis, translations):
return pixel_translations[0], propagations[0]
else:
return pixel_translations, propagations
def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj=False):
"""Returns a stack of exit waves without accounting for subpixel shifts
@@ -229,15 +229,15 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj
dimension as the translation index and the final dimensions
corresponding to the detector. The exit waves are calculated by
shifting the probe by the rounded value of the translation
If multiple_modes is set to False, any additional dimensions in the
ptycho_2D_round function will be assumed to correspond to the translation
index. If multiple_modes is set to true, the (-4th) dimension of the probe
will always be assumed to be defining a set of (P) incoherently mixing
modes to be broadcast all translation indices. If any additional dimensions
closer to the start exist, they will be assumed to be translation indices
Parameters
----------
probe : torch.Tensor
@@ -251,7 +251,7 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj
Returns
-------
exit_waves : torch.Tensor
exit_waves : torch.Tensor
An (N)x(P)xMxL tensor of the calculated exit waves
"""
@@ -260,9 +260,9 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj
translations = translations[None,:]
single_translation = True
integer_translations = t.round(translations).to(dtype=t.int32)
if upsample_obj:
selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2]//2,
tr[1]:tr[1]+probe.shape[-1]//2]
@@ -292,7 +292,7 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj
def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
"""Returns a stack of exit waves accounting for subpixel shifts
This function returns a collection of exit waves, with the first
dimension as the translation index and the final dimensions
corresponding to the detector. The exit waves are calculated by
@@ -322,7 +322,7 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
if translations.dim() == 1:
translations = translations[None,:]
single_translation = True
# Separate the translations into a part that chooses the window
# And a part that defines the windowing function
integer_translations = t.floor(translations)
@@ -342,15 +342,15 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
sel01 = t.cat((probe[:,-1:],probe[:,:-1]),dim=1)
sel10 = t.cat((probe[-1:,:],probe[:-1,:]),dim=0)
sel11 = t.cat((sel01[-1:,:],sel01[:-1,:]),dim=0)
selection = sel00 * (1-sp[0])*(1-sp[1]) + \
sel10 * sp[0]*(1-sp[1]) + \
sel01 * (1-sp[0])*sp[1] + \
sel11 * sp[0]*sp[1]
obj_slice = obj[tr[0]:tr[0]+probe.shape[0],
tr[1]:tr[1]+probe.shape[1]]
exit_waves.append(selection * obj_slice)
else:
for tr, sp in zip(integer_translations,
@@ -359,16 +359,16 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
# Here we subpixel shift the object by (-i,-j) after
# slicing out the correct translation of the probe
#
sel00 = obj[tr[0]:tr[0]+probe.shape[0],
tr[1]:tr[1]+probe.shape[1]]
sel01 = obj[tr[0]:tr[0]+probe.shape[0],
tr[1]+1:tr[1]+1+probe.shape[1]]
sel10 = obj[tr[0]+1:tr[0]+1+probe.shape[0],
tr[1]:tr[1]+probe.shape[1]]
sel11 = obj[tr[0]+1:tr[0]+1+probe.shape[0],
tr[1]+1:tr[1]+1+probe.shape[1]]
@@ -387,7 +387,7 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multiple_modes=True, polarized=False, polarizer=None, analyzer=None):
"""Returns a stack of exit waves accounting for subpixel shifts
This function returns a collection of exit waves, with the first
dimension as the translation index and the final dimensions
corresponding to the detector. The exit waves are calculated by
@@ -427,7 +427,7 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
if translations.dim() == 1:
translations = translations[None,:]
single_translation = True
# Separate the translations into a part that chooses the window
# And a part that defines the windowing function
integer_translations = t.floor(translations)
@@ -480,7 +480,7 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
else:
raise NotImplementedError('Object shift not yet implemented')
print('ptyvho 2d sinc', output.shape)
if single_translation:
return output[0]
else:
@@ -489,7 +489,7 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, padding=10):
"""Returns a stack of exit waves accounting for subpixel shifts
This function returns a collection of exit waves, with the first
dimension as the translation index and the final dimensions
corresponding to the detector. The exit waves are calculated by
@@ -505,7 +505,7 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
on the input wavefield, and the first two indexes index differences from
that pixel. It is easier to interpret the resulting matrix though if the
latter two indices index locations in the output plane. NOTE: I believe
this change has now been made
this change has now been made
Parameters
----------
@@ -529,17 +529,17 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
if translations.dim() == 1:
translations = translations[None,:]
single_translation = True
# Separate the translations into a part that chooses the window
# And a part that defines the windowing function
integer_translations = t.floor(translations)
subpixel_translations = translations - integer_translations
integer_translations = integer_translations.to(dtype=t.int32)
exit_waves = []
B = s_matrix.shape[0]//2
if shift_probe:
i = t.arange(probe.shape[-2]) - probe.shape[-2]//2
j = t.arange(probe.shape[-1]) - probe.shape[-1]//2
@@ -548,14 +548,14 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
J = 2 * np.pi * J.to(t.float32) / probe.shape[-1]
I = I.to(dtype=probe.dtype,device=probe.device)
J = J.to(dtype=probe.dtype,device=probe.device)
for tr, sp in zip(integer_translations,
subpixel_translations):
fft_probe = t.fft.fftshift(t.fft.fft2(probe), dim=(-1,-2))
shifted_fft_probe = fft_probe * t.exp(1j*(-sp[0]*I - sp[1]*J))
shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe,
dim=(-1,-2)))
s_matrix_slice = s_matrix[:,:,tr[0]:tr[0]+probe.shape[-2]+2*B,
tr[1]:tr[1]+probe.shape[-1]+2*B]
@@ -564,14 +564,14 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
device=s_matrix_slice.device,
dtype=s_matrix_slice.dtype)
for i in range(s_matrix.shape[0]):
for j in range(s_matrix.shape[1]):
output [i:i+probe.shape[-2],j:j+probe.shape[-1]] += \
shifted_probe * s_matrix_slice[i,j,i:i+probe.shape[-2],j:j+probe.shape[-1]]
exit_waves.append(output)
else:
raise NotImplementedError('Object shift not yet implemented')
@@ -579,11 +579,11 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
return exit_waves[0]
else:
return t.stack(exit_waves)
def RPI_interaction(probe, obj):
"""Returns an exit wave from a high-res probe and a low-res obj
In this interaction, the probe and object arrays are assumed to cover
the same physical region of space, but with the probe array sampling that
region of space more finely. Thus, to do the interaction, the object
@@ -593,7 +593,7 @@ def RPI_interaction(probe, obj):
method and is not commonly used elsewhere.
This also works with object functions that have an extra first dimension
for an incoherently mixing model.
for an incoherently mixing model.
Parameters
@@ -610,7 +610,7 @@ def RPI_interaction(probe, obj):
"""
# TODO: The upsampling only works for arrays of even dimension!
# The far-field propagator is just a 2D FFT but with an fftshift
fftobj = propagators.far_field(obj)
# We calculate the padding that we need to do the upsampling
@@ -618,7 +618,7 @@ def RPI_interaction(probe, obj):
pad0r = probe.shape[-2] - obj.shape[-2] - pad0l
pad1l = (probe.shape[-1] - obj.shape[-1])//2
pad1r = probe.shape[-1] - obj.shape[-1] - pad1l
if obj.dim() == 2:
fftobj = t.nn.functional.pad(fftobj, (pad1l, pad1r, pad0l, pad0r))
elif obj.dim() == 3:
@@ -626,7 +626,7 @@ def RPI_interaction(probe, obj):
fftobj, (pad1l, pad1r, pad0l, pad0r, 0,0))
else:
raise NotImplementedError('RPI interaction with obj of dimension higher than 4 (including complex dimension) is not supported.')
# Again, just an inverse FFT but with an fftshift
upsampled_obj = propagators.inverse_far_field(fftobj)
CDTools/tools/plotting/plotting.py
+45 -41
View File
@@ -17,7 +17,11 @@ from matplotlib import ticker, patheffects
__all__ = ['colorize', 'plot_amplitude', 'plot_phase',
'plot_colorized', 'plot_translations', 'get_units_factor',
'plot_nanomap', 'plot_real', 'plot_imag',
'plot_nanomap_with_images']
'plot_nanomap_with_images',
'polarized_plot_component_amplitudes',
'polarized_plot_phase_ret',
'polarized_plot_global_phases',
'polarized_plot_ellipses']
def colorize(z):
@@ -84,7 +88,7 @@ def get_units_factor(units):
def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label=None, **kwargs):
"""Plots an image with a colorbar and on an appropriate spatial grid
If a figure is given explicitly, it will clear that existing figure and
plot over it. Otherwise, it will generate a new figure.
@@ -94,7 +98,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
Finally, if a function is passed to the plot_func argument, this function
will be called on each slice of data before it is plotted. This is used
internally to enable the plot_real, plot_image, plot_phase, etc. functions.
Parameters
----------
@@ -120,7 +124,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
# convert to numpy
if isinstance(im, t.Tensor):
# If final dimension is 2, assume it is a complex array. If not,
@@ -142,7 +146,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
title = plt.gca().get_title()
fig.clear()
# If im only has two dimensions, this reshape will add a leading
# dimension, and update will be called on index 0. If it has 3 or more
# dimensions, then all the leading dimensions will be compressed into
@@ -151,9 +155,9 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
reshaped_im = im.reshape(-1,s[-2],s[-1])
num_images = reshaped_im.shape[0]
fig.plot_idx = idx % num_images
to_plot = plot_func(reshaped_im[fig.plot_idx])
#Plot in a basis if it exists, otherwise dont
if basis is not None:
if isinstance(basis,t.Tensor):
@@ -181,7 +185,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
plt.title(title)
if len(im.shape) >= 3:
@@ -194,14 +198,14 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
result = make_plot(0)
update = make_plot
def on_action(event):
if not hasattr(event, 'button'):
event.button = None
if not hasattr(event, 'key'):
event.key = None
if event.key == 'up' or event.button == 'up':
update(fig.plot_idx - 1)
elif event.key == 'down' or event.button == 'down':
@@ -217,9 +221,9 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
fig.my_callbacks = []
fig.my_callbacks.append(fig.canvas.mpl_connect('key_press_event',on_action))
fig.my_callbacks.append(fig.canvas.mpl_connect('scroll_event',on_action))
return result
def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Real Part (a.u.)', **kwargs):
"""Plots the real part of a complex array with dimensions NxM
@@ -256,7 +260,7 @@ def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, cmap=cmap, cmap_label=cmap_label,
**kwargs)
def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Imaginary Part (a.u.)', **kwargs):
@@ -304,7 +308,7 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis',
If a basis is explicitly passed, the image will be plotted in real-space
coordinates.
Parameters
----------
im : array
@@ -520,12 +524,12 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
def plot_nanomap_with_images(translations, get_image_func, values=None, mask=None, basis=None, fig=None, nanomap_units='$\\mu$m', image_units='$\\mu$m', convention='probe', image_title='Image', image_colorbar_title='Image Amplitude', nanomap_colorbar_title='Integrated Intensity', cmap='viridis', **kwargs):
"""Plots a nanomap, with an image or stack of images for each point
In many situations, ptychography data or the output of ptychography
reconstructions is formatted as a set of images associated with various
points in real space. This function is designed to allow for browsing
through this kind of data, by making it possible to visualize a
"""
# This should pull heavily from the dataset.inspect function
@@ -558,11 +562,11 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
s0 = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch
s0 /= 4 # A rough value to make the size work out
s = np.ones(translations.shape[0]) * s0
s[idx] *= 4
return s
def update_colorbar(im):
#
# This solves the problem of the colorbar being changed
@@ -571,29 +575,29 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
if hasattr(im, 'norecurse') and im.norecurse:
im.norecurse=False
return
im.norecurse=True
# This is needed to update the colorbar
# only change limits if array contains multiple values
if np.min(im.get_array()) != np.max(im.get_array()):
im.set_clim(vmin=np.min(im.get_array()),
vmax=np.max(im.get_array()))
#
# The meatiest part of this program, here we just go through and
# set up the plot how we want it
#
# First we set up the left-hand plot, which shows an overview map
axes[0].set_title('Relative Displacement Map')
translations = translations.detach().cpu().numpy()
if convention.lower() != 'probe':
translations = translations * -1
s = calculate_sizes(0)
nanomap_units_factor = get_units_factor(nanomap_units)
nanomap = axes[0].scatter(nanomap_units_factor * translations[:,0],
nanomap_units_factor * translations[:,1],
@@ -614,7 +618,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# where the colorbar should have been to avoid stretching the
# nanomap plot, while still not showing the (now useless) colorbar.
cb1.remove()
# Now we set up the second plot, which shows the individual
# diffraction patterns
axes[1].set_title(image_title)
@@ -634,7 +638,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# This fails if the basis is not rectangular
basis_norm = np.linalg.norm(np_basis, axis = 0)
basis_norm = basis_norm * get_units_factor(image_units)
extent = [0, example_im.shape[-1]*basis_norm[1], 0,
example_im.shape[-2]*basis_norm[0]]
else:
@@ -655,9 +659,9 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
axes[1].text_box.set_path_effects(
[patheffects.Stroke(linewidth=2, foreground='black'),
patheffects.Normal()])
meas = axes[1].imshow(im, extent=extent, cmap=cmap)
cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',
format='%.2e',
ticks=ticker.LinearLocator(numticks=5),
@@ -665,7 +669,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
cb2.ax.tick_params(labelrotation=20)
cb2.ax.set_title(image_colorbar_title, size="medium", pad=5)
cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas))
# This function handles all the updating, except for moving the
# slider value. This is done because the slider widget is
# ultimately responsible for triggering an update, so all other
@@ -675,7 +679,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# We have to explicitly make it an integer because the slider will
# output floats (even if they are still integer-valued)
idx = int(idx)
# Get the new data for this index
im = get_image_func(idx)
if len(im.shape) >= 3:
@@ -686,22 +690,22 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
axes[1].image_idx = im_idx
axes[1].text_box.set_text(str(im_idx))
im = im.reshape(-1,im.shape[-2],im.shape[-1])[im_idx]
# Now we resize the nanomap to show the new selection
axes[0].collections[0].set_sizes(calculate_sizes(idx))
# And we update the data in the image as well
ax_im = axes[1].images[-1]
ax_im.set_data(im)
update_colorbar(ax_im)
#
# Now we define the functions to handle various kinds of events
# that can be thrown our way
#
# We start by creating the slider here, so it can be used
# by the update hooks.
slider = Slider(axslider, 'Image #', 0, translations.shape[0]-1, valstep=1, valfmt="%d")
@@ -715,7 +719,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# while the mouse is within the image display
im = im.reshape(-1,im.shape[-2],im.shape[-1])
im_idx = axes[1].image_idx
if event.key == 'up' or event.button == 'up' \
or event.key == 'left':
im_idx = (im_idx - 1) % im.shape[0]
@@ -733,7 +737,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
event.button = None
if not hasattr(event, 'key'):
event.key = None
if event.key == 'up' or event.button == 'up' or event.key == 'left':
idx = slider.val - 1
elif event.key == 'down' or event.button == 'down' or event.key == 'right':
@@ -742,7 +746,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# This prevents errors from being thrown on irrelevant key
# or mouse input
return
# Handle the wraparound and trigger the update
idx = int(idx) % translations.shape[0]
slider.set_val(idx)
@@ -753,16 +757,16 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
# for example, scroll events that happen over the nanomap
if event.mouseevent.button == 1:
slider.set_val(event.ind[0])
# Here we connect the various update functions
cid1 = fig.canvas.mpl_connect('pick_event',on_pick)
cid2 = fig.canvas.mpl_connect('key_press_event',on_action)
cid3 = fig.canvas.mpl_connect('scroll_event',on_action)
# It's so dumb that matplotlib doesn't automatically track this for you
fig.nanomap_cids = [cid1,cid2,cid3]
fig.nanomap_cids = [cid1,cid2,cid3]
slider.on_changed(update)
# Throw an extra update into the mix just to get rid of any things
# (like the nanomap dot sizes) that otherwise would change on the
# first update
CDTools/tools/polarization/polarization.py
+30 -25
View File
@@ -14,7 +14,7 @@ __all__ = ['apply_linear_polarizer',
'apply_circular_polarizer',
'apply_jones_matrix',
'generate_linear_polarizer',
'generate_phase_retarder']
'generate_birefringent_obj']
# Abe - split these into two functions
@@ -36,10 +36,10 @@ def generate_linear_polarizer(pol_angle):
cd = t.stack((c, d), dim=-1)
jones_matrices = t.stack((ab, cd), dim=-2)
if single_angle:
return jones_matrices[0].to(dtype=t.cfloat)
return jones_matrices[0].to(dtype=t.cfloat)
else:
return jones_matrices.to(dtype=t.cfloat)
def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True):
"""
@@ -56,7 +56,7 @@ def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True
Returns:
--------
linearly polarized probe: t.Tensor
(N)(P)x2x1xMxL
(N)(P)x2x1xMxL
"""
jones_matrices = generate_linear_polarizer(polarizer)
return apply_jones_matrix(probe, jones_matrices, transpose=transpose, multiple_modes=multiple_modes)
@@ -75,19 +75,19 @@ def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True)
probe: t.Tensor
A (N)(P)x2xMxL tensor representing the probe
jones_matrix: t.tensor
(N)x2x2x(M)x(L)
(N)x2x2x(M)x(L)
Returns:
--------
a probe with the jones matrix applied: t.Tensor
(N)(P)x2xMxL
(N)(P)x2xMxL
"""
if transpose:
if jones_matrix.dim() < 4:
jones_matrix = jones_matrix[..., None, None]
if multiple_modes:
jones_matrix = jones_matrix.unsqueeze(-5)
jones_matrix = jones_matrix.unsqueeze(-5)
probe = probe[..., None, :, :]
# if jones matrices do not differ from pattern to pattern
if probe.dim() > jones_matrix.dim():
@@ -96,19 +96,19 @@ def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True)
elif jones_matrix.dim() > probe.dim():
probe = probe.unsqueeze(0)
# print('apply jonesmatrix: probe', probe.shape, 'matrix:', jones_matrix)
jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4)
jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4)
probe = probe.transpose(-1, -3).transpose(-2, -4)
output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3)
else:
raise NotImplementedError
return output
def apply_phase_retardance(probe, phase_shift, multiple_modes=True):
"""
Shifts the y-component of the field wrt the x-component by a given phase shift
Shifts the y-component of the field wrt the x-component by a given phase shift
Parameters:
----------
@@ -120,7 +120,7 @@ def apply_phase_retardance(probe, phase_shift, multiple_modes=True):
Returns:
--------
probe: t.Tensor
(...)x2x1xMxL
(...)x2x1xMxL
"""
theta = t.as_tensor(phase_shift, dtype=t.float32)
theta = t.deg2rad(theta)
@@ -140,11 +140,11 @@ def apply_circular_polarizer(probe, left_polarized=True, multiple_modes=True):
A (...)x2xMxL tensor representing the probe
left_polarizd: bool
True for the left-polarization, False for the right
Returns:
--------
circularly polarized probe: t.Tensor
(...)x2xMxL
(...)x2xMxL
"""
probe = probe.to(dtype=t.cfloat)
if left_polarized:
@@ -166,7 +166,7 @@ def apply_quarter_wave_plate(probe, fast_axis_angle, multiple_modes=True):
Returns:
--------
polarized probe: t.Tensor
(...)x2x1xMxL
(...)x2x1xMxL
"""
probe = probe.to(dtype=t.cfloat)
theta = math.radians(fast_axis_angle)
@@ -174,7 +174,7 @@ def apply_quarter_wave_plate(probe, fast_axis_angle, multiple_modes=True):
jones_matrix = exponent* t.tensor([[(cos(theta))**2 + 1j * (sin(theta))**2, (1 - 1j) * sin(theta) * cos(theta)], [(1 - 1j) * sin(theta) * cos(theta), (sin(theta))**2 + 1j * (cos(theta))**2]]).to(dtype=t.cfloat)
out = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes)
return out
return out
def apply_half_wave_plate(probe, fast_axis_angle, multiple_modes=True):
"""
@@ -188,7 +188,7 @@ def apply_half_wave_plate(probe, fast_axis_angle, multiple_modes=True):
Returns:
--------
polarized probe: t.Tensor
(...)x2x1xMxL
(...)x2x1xMxL
"""
probe = probe.to(dtype=t.cfloat)
theta = math.radians(fast_axis_angle)
@@ -196,19 +196,24 @@ def apply_half_wave_plate(probe, fast_axis_angle, multiple_modes=True):
jones_matrix = exponent * t.tensor([[(cos(theta))**2 - (sin(theta))**2, 2 * sin(theta) * cos(theta)], [2 * sin(theta) * cos(theta), (sin(theta))**2 - (cos(theta))**2]]).to(dtype=t.cfloat)
out = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes)
return out
def generate_phase_retarder(fast_axis=0, phase=0):
phase = t.as_tensor(phase).to(dtype=t.float32)
phase = t.deg2rad(phase)
def coord_rot(angle):
return out
def generate_birefringent_obj(fast_axis=90, phase_ret=10, atten_fast=1, atten_ret=1, global_phase=0):
def to_rad(angle):
angle = t.as_tensor(angle, dtype=t.float32)
angle = t.deg2rad(angle)
return angle
fast_axis = to_rad(fast_axis)
phase_ret = to_rad(phase_ret)
global_phase = to_rad(global_phase)
def coord_rot(angle):
a = t.stack((t.cos(angle), t.sin(angle)), dim=-1)
b = t.stack((-t.sin(angle), t.cos(angle)), dim=-1)
return t.stack((a, b), dim=-2).to(dtype=t.cfloat)
r1 = coord_rot(-fast_axis)
r2 = coord_rot(fast_axis)
p = t.as_tensor([[1, 0], [0, t.exp(phase*1j)]], dtype=t.cfloat)
return t.matmul(r1, t.matmul(p, r2))
p = t.exp(global_phase * 1j) * t.as_tensor([[atten_fast, 0], [0, atten_ret * t.exp(phase_ret*1j)]], dtype=t.cfloat)
return t.matmul(r1, t.matmul(p, r2))
temp_tests/simulated_dataset.py
+107 -62
View File
@@ -1,7 +1,8 @@
import numpy as np
import torch as t
from scipy import misc
from CDTools.models import PolarizedFancyPtycho
#from CDTools.datasets import Polarized2DDataset
from CDTools.datasets import PolarizedPtycho2DDataset
import CDTools
from CDTools.tools import polarization
from CDTools import tools
@@ -11,77 +12,121 @@ from PIL import Image
# upolad 4 different images representing 4 components of the object
# and 2 gaaussian functionas corresponding to the probe components
a = np.asarray(Image.open('a.jpg'))
b = np.asarray(Image.open('b.jpg'))
c = np.asarray(Image.open('c.jpg'))
d = np.asarray(Image.open('d.jpg'))
f = misc.ascent()
x , y = np.shape(f)
aa = f[:x//2, :y//2]
bb = f[:x//2, -y//2:]
cc = f[-x//2:, :y//2]
dd = f[-x//2:, -y//2:]
#a = np.dot(a[..., :3], [.3, 6., .1])
print(1)
def simulate_dataset(probe_size, obj_size, num_patt):
translations = []
xs, ys = np.mgrid[:num_patt, :num_patt]
for x, y in zip(xs, ys):
translations.append((x*10e-3, y*10e-3))
translations = t.as_tensor(translations, dtype=t.float32)
a = t.as_tensor(a, dtype=t.cfloat)
a = t.tensordot(a, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1],[0]))[:obj_size, :obj_size]
probe = tools.initializers.gaussian(np.array([probe_size, probe_size]), 50)
wavefields = tools.interactions.ptycho_2D_sinc(probe, obj, translations)
patterns = tools.propagators.far_field(wavefront)
patterns = np(patterns)
translations = np(t.cat((translations, t.zeros(num_patt)), dim=-1))
# needs to be stored as a cxi file
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi('simulated_dataset.cxi')
dataset.detector_geometry = None
def simulate polarized_datset(probe_size, obj_size, num_patt):
def simulate_polarized_dataset(probe_size, obj_size, num_patt, a, b, c, d):
a, b, c, d = t.as_tensor(a, dtype=t.cfloat), t.as_tensor(b, dtype=t.cfloat), t.as_tensor(c, dtype=t.cfloat), t.as_tensor(d, dtype=t.cfloat)
a = t.tensordot(a, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1],[0]))[:obj_size, :obj_size]
b = t.tensordot(b, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
c = t.tensordot(c, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
d = t.tensordot(d, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
# a = t.tensordot(a, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1],[0]))[:obj_size, :obj_size]
# b = t.tensordot(b, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
# c = t.tensordot(c, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
# d = t.tensordot(d, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
translations = []
xs, ys = np.mgrid[:num_patt, :num_patt]
xs, ys = np.ravel(xs), np.ravel(ys)
for x, y in zip(xs, ys):
translations.append((x*10e-3, y*10e-3))
for j in range(9):
translations.append((5*x, 5*y))
translations = t.as_tensor(translations, dtype=t.float32)
obj = t.stack((t.stack((a, c), dim=0), t.stack((b, d), dim=0)), dim=-3)
probe = tools.initializers.gaussian(np.array([probe_size, probe_size]), 50)
obj = t.stack((t.stack((a, b), dim=-3), t.stack((c, d), dim=-3)), dim=-4)
# for i, j in zip([0, 0, 1, 1], [0, 1, 0, 1]):
# plt.imshow(np.real(obj[i, j, ...]))
# plt.show()
probe = tools.initializers.gaussian(np.array([probe_size, probe_size]), np.array((2, 2)))
probe = t.stack((probe, probe), dim=-3)
probe = polarization.apply_circular_polarizer(probe)
probe = polarization.apply_circular_polarizer(probe, multiple_modes=False)
polarizers = polarization.generate_linear_polarizer(t.tensor([0, 45, 90]))
num_transl = len(translations)
pol_probes = [polarization.apply_jones_matrix(probe, polarizers[i], multiple_modes=False) for i in range(3)]
analyzers = t.stack(([polarizers[i % 3] for i in range(num_transl)]), dim=0)
analyzer = t.tensor([(i % 3) * 45 for i in range(num_transl)])
polarizer = t.tensor([(i // 3) % 3 * 45 for i in range(num_transl)])
polarized_probes = t.stack(([pol_probes[(i // 3) % 3] for i in range(num_transl)]), dim=0)
polarized_wavefields = t.stack([tools.interactions.ptycho_2D_sinc(polarized_probes[i], obj, translations[i], polarized=True, multiple_modes=False) for i in range(num_transl)])
wavefields = tools.propagators.far_field(polarized_wavefields)
wavefields = polarization.apply_jones_matrix(wavefields, analyzers, multiple_modes=False)
patterns = t.abs(wavefields[:, 0, :, :])**2 + t.abs(wavefields[:, 1, :, :])**2
detector_basis = t.transpose(t.tensor([[0, -4.8e-6, 0], [-4.8e-6, 0, 0]]), 0, 1)
det_shape = t.Size((obj_size, obj_size))
wavelength = 532e-9
real_basis = tools.initializers.exit_wave_geometry(detector_basis, det_shape, wavelength, 2.5e-2)[0]
# print('f', type(real_basis), translations.shape)
real_translations = tools.interactions.pixel_to_translations(real_basis, translations)
# print(type(real_translations))
patterns = patterns.numpy()
detector_geometry = {
'corner': np.array([probe_size*4.8e-6/2, probe_size*4.8e-6/2, 2.5e-2]),
'basis': np.array([[0, -4.8e-6, 0], [-4.8e-6, 0, 0]]).transpose(),
'distance': 2.5e-2
}
selections = tools.interactions.ptycho_2D_sinc(t.ones(2, probe_size, probe_size).to(dtype=t.cfloat), obj, translations, polarized=True)
polarizers = [polarization.generate_linear_polarizer(i * 45) for i in range(3)]
pol_probes = [polarization.apply_jones_matrix(probe, polarizers[i]) for i in range(3)]
return PolarizedPtycho2DDataset(real_translations, polarizer, analyzer, patterns,
axes=("x", "y"), detector_geometry=detector_geometry, wavelength=wavelength)
analyzer = t.stack(([polaryzers[i % 3] for i in range(num_patt)]), dim=0)
# probes = t.stack(([probes[i // 3] for i in num_patt]), dim=0)
wavefields = t.as_tensor([tools.interactions.ptycho_2D_sinc(pol_probes[i], obj, translations, polarized=True) for i in range(3)]).to(dtype=t.cfloat)
# print(1)
dataset = simulate_polarized_dataset(100, 500, 10, aa, bb, cc, dd)
# print(2)
dataset.inspect()
# print(3)
plt.show()
model = PolarizedFancyPtycho.from_dataset(dataset, propagation_distance=1e-3)
model.inspect()
# plt.show()
model.compare(dataset)
plt.show()
wf = t.empty(1, 2, obj_size, obj_size)
for i in range(num_patt):
for j in range(3):
pol_channel = t.stack(([wavefileds[j] for k in range(3)]), dim=0)
wf = t.cat((wf, pol_channel), dim=0)
pol_wavefieds = polarization.apply_jones_matrix(wf, analyzer)
# for loss in model.Adam_optimize(400, dataset, batch_size=5, lr=0.002, schedule=True):
# # And we liveplot the updates to the model as they happen
# print(model.report())
# model.inspect(dataset)
# model.save_figures(prefix='simulated', extension='png')
# res = model.save_results(dataset)
# np.save('simulated_dataset.npy', res)
patterns = tools.propagators.far_field(pol_wavefieds)
dataset = simulate_polarized_dataset(50, 100, 10, aa, bb, cc, dd)
translations = np(t.cat((translations, t.zeros(num_patt)), dim=-1))
patterns = np(patterns)
dataset.detector_geometry = None
# needs to be stored in a cxi file
dataset = CDTools.datasets.FancyPtycho2DDataset.from_cxi('polarized_simulated_dataset.cxi')
dataset.inspect()
model = tools.models.PolarizedFancyPtycho.from_dataset(dataset)
res = np.load('simulated_dataset.npy', allow_pickle=True)
res = res[()]
print(type(res))
Ws = t.ones(len(dataset))
ewg = CDTools.tools.initializers.exit_wave_geometry
probe_basis, probe_shape, det_slice = ewg(res['basis'],
dataset[0][1].shape,
dataset.wavelength,
dataset.detector_geometry['distance'],
center=None,
padding=0)
print('object', res['obj'].shape)
obj = res['obj'][..., 200:275, 200:275]
# print('obj', obj)
a = obj[0, 0, :, :]
b = obj[0, 1, :, :]
c = obj[1, 0, :, :]
d = obj[1, 1, :, :]
for i in [a, b, c, d]:
plt.imshow(np.real(i))
# plt.imshow(np.log(np))
plt.colorbar()
plt.show()
print(np.allclose(np.real(a), np.real(b)))
print('a', a[..., 10:20, 10:20])
print('b', b[..., 10:20, 10:20])
models = [CDTools.models.FancyPtycho(dataset.wavelength, dataset.detector_geometry, probe_basis,
res['probe'], component, surface_normal=t.tensor([0., 0., 1.], dtype=t.float32),
min_translation=t.tensor([0, 0], dtype=t.float32),
background=t.tensor(res['background']), translation_offsets=None, mask=None,
weights=Ws, translation_scale=1, saturation=None,
probe_support=None, oversampling=1,
loss='amplitude mse', units='um') for component in [a, b, c, d]]
# print('model is created')
# for model in models:
# model.inspect()
# plt.show()