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

Merging a whole buncha work back into master

Browse Source
This commit is contained in:
Abe Levitan
2022-04-18 15:51:45 -07:00
parent 1b5ef2271c 20a1479859
commit f3cee949e0
22 changed files with 3382 additions and 771 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/datasets/polarized_ptycho_2d_dataset.py
+6 -3
View File
@@ -115,7 +115,7 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset):
# It sucks that I can't reuse the base factory method here,
# perhaps there is a way but I couldn't figure it out.
@classmethod
def from_cxi(cls, cxi_file):
def from_cxi(cls, cxi_file, cut_zeros=True):
"""Generates a new PolarizedPtycho2DDataset from a .cxi file directly
This generates a new PolarizedPtycho2DDataset from a .cxi file storing
@@ -125,19 +125,22 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset):
----------
file : str, pathlib.Path, or h5py.File
The .cxi file to load from
cut_zeros : bool
Default True, whether to set all negative data to zero
Returns
-------
dataset : PolarizedPtycho2DDataset
The constructed dataset object
"""
# If a bare string is passed
if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path):
with h5py.File(cxi_file, 'r') as f:
return cls.from_cxi(f)
return cls.from_cxi(f, cut_zeros=cut_zeros)
# Generate a base dataset
dataset = Ptycho2DDataset.from_cxi(cxi_file)
dataset = Ptycho2DDataset.from_cxi(cxi_file, cut_zeros=cut_zeros)
# Mutate the class to this subclass (PolarizedPtycho2DDataset)
dataset.__class__ = cls
CDTools/datasets/ptycho_2d_dataset.py
+26 -11
View File
@@ -19,7 +19,8 @@ class Ptycho2DDataset(CDataset):
It should save and load files compatible with most reconstruction
programs, although it is only tested against SHARP.
"""
def __init__(self, translations, patterns, axes=None, *args, **kwargs):
def __init__(self, translations, patterns, intensities=None,
axes=None, *args, **kwargs):
"""The __init__ function allows construction from python objects.
The detector_geometry dictionary is defined to have the
@@ -52,23 +53,27 @@ class Ptycho2DDataset(CDataset):
background : array
An initial guess for the not-previously-subtracted
detector background
intensities : array
A list of measured shot-to-shot intensities
"""
super(Ptycho2DDataset,self).__init__(*args, **kwargs)
self.axes = copy(axes)
self.translations = t.tensor(translations, dtype=t.float32)
self.translations = t.tensor(translations)
self.patterns = t.as_tensor(patterns, dtype=t.float32)
if self.patterns.dtype == t.float64:
raise NotImplementedError('64-bit floats are not supported and precision will not be retained in reconstructions! Please explicitly convert your data to 32-bit or submit a pull request')
self.patterns = t.as_tensor(patterns)
if self.mask is None:
self.mask = t.ones(self.patterns.shape[-2:]).to(dtype=t.bool)
self.mask.masked_fill_(t.isnan(t.sum(self.patterns,dim=(0,))),0)
self.patterns.masked_fill_(t.isnan(self.patterns),0)
if intensities is not None:
self.intensities = t.as_tensor(intensities, dtype=t.float32)
else:
self.intensities = None
def __len__(self):
return self.patterns.shape[0]
@@ -118,7 +123,7 @@ class Ptycho2DDataset(CDataset):
# It sucks that I can't reuse the base factory method here,
# perhaps there is a way but I couldn't figure it out.
@classmethod
def from_cxi(cls, cxi_file):
def from_cxi(cls, cxi_file, cut_zeros=True):
"""Generates a new Ptycho2DDataset from a .cxi file directly
This generates a new Ptycho2DDataset from a .cxi file storing
@@ -128,6 +133,8 @@ class Ptycho2DDataset(CDataset):
----------
file : str, pathlib.Path, or h5py.File
The .cxi file to load from
cut_zeros : bool
Default True, whether to set all negative data to zero
Returns
-------
@@ -145,11 +152,10 @@ class Ptycho2DDataset(CDataset):
dataset.__class__ = cls
# Load the data that is only relevant for this class
patterns, axes = cdtdata.get_data(cxi_file)
patterns, axes = cdtdata.get_data(cxi_file, cut_zeros=cut_zeros)
translations = cdtdata.get_ptycho_translations(cxi_file)
# And now re-do the stuff from __init__
dataset.translations = t.tensor(translations, dtype=t.float32)
dataset.patterns = t.as_tensor(patterns)
if dataset.patterns.dtype == t.float64:
raise NotImplementedError('64-bit floats are not supported and precision will not be retained in reconstructions! Please explicitly convert your data to 32-bit or submit a pull request')
@@ -158,6 +164,12 @@ class Ptycho2DDataset(CDataset):
if dataset.mask is None:
dataset.mask = t.ones(dataset.patterns.shape[-2:]).to(dtype=t.bool)
try:
intensities = cdtdata.get_shot_to_shot_info(cxi_file, 'intensities')
dataset.intensities = t.as_tensor(intensities, dtype=t.float32)
except KeyError:
dataset.intensities = None
return dataset
@@ -187,6 +199,9 @@ class Ptycho2DDataset(CDataset):
cdtdata.add_data(cxi_file, self.patterns)
cdtdata.add_ptycho_translations(cxi_file, self.translations)
if hasattr(self, 'intensities') and self.intensities is not None:
cdtdata.add_shot_to_shot_info(cxi_file, self.intensities, 'intensities')
def inspect(self, logarithmic=True, units='um', log_offset=1):
"""Launches an interactive plot for perusing the data
@@ -234,5 +249,5 @@ class Ptycho2DDataset(CDataset):
else:
cbar_title = 'Diffraction Intensity'
plotting.plot_nanomap_with_images(self.translations.detach().cpu(), get_images, values=nanomap_values, nanomap_units=units, image_title='Diffraction Pattern', image_colorbar_title=cbar_title)
return plotting.plot_nanomap_with_images(self.translations.detach().cpu(), get_images, values=nanomap_values, nanomap_units=units, image_title='Diffraction Pattern', image_colorbar_title=cbar_title)
CDTools/models/__init__.py
+4 -1
View File
@@ -22,7 +22,7 @@ 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', 'PolarizedFancyPtycho', 'Bragg2DPtycho', 'Multislice2DPtycho', 'RPI']
__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'PolarizedFancyPtycho', 'Bragg2DPtycho', 'Multislice2DPtycho', 'RPI', 'TimeResolvedPtychoCalibration', 'TimeResolvedRPI']
from CDTools.models.base import CDIModel
from CDTools.models.simple_ptycho import SimplePtycho
@@ -31,6 +31,9 @@ from CDTools.models.polarized_fancy_ptycho import PolarizedFancyPtycho
from CDTools.models.bragg_2d_ptycho import Bragg2DPtycho
from CDTools.models.multislice_2d_ptycho import Multislice2DPtycho
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
# Still needs to be updated for the new complex numbers
#from CDTools.models.s_matrix_ptycho import SMatrixPtycho
CDTools/models/base.py
+76 -61
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
----------
@@ -55,8 +55,10 @@ class CDIModel(t.nn.Module):
def __init__(self):
super(CDIModel,self).__init__()
self.loss_train = []
self.iteration_count = 0
def from_dataset(self, dataset):
raise NotImplementedError()
@@ -79,11 +81,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 +101,7 @@ class CDIModel(t.nn.Module):
def simulate_to_dataset(self, args_list):
raise NotImplementedError()
def save_results(self):
raise NotImplementedError()
@@ -107,10 +109,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 +137,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 +146,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 +154,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 +170,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 +180,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)
@@ -197,7 +199,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)
@@ -205,10 +207,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.')
@@ -232,7 +234,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
@@ -266,7 +268,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)
@@ -289,17 +291,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):
calculation_width=10, line_search_fn=None):
"""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
@@ -319,14 +321,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))
@@ -334,10 +336,11 @@ class CDIModel(t.nn.Module):
# Define the optimizer
optimizer = t.optim.LBFGS(self.parameters(),
lr = lr, history_size=history_size)
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,
thread=thread,
@@ -349,7 +352,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.
@@ -381,7 +384,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,
@@ -417,28 +420,28 @@ class CDIModel(t.nn.Module):
self.latest_iteration_time + str(self.loss_train[-1])
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
@@ -447,9 +450,20 @@ 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
@@ -475,7 +489,7 @@ class CDIModel(t.nn.Module):
plotter = plots[1]
if figs is None:
fig = plt.figure()
fig = plt.figure()
self.figs.append(fig)
else:
fig = figs[idx]
@@ -483,13 +497,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
@@ -497,15 +511,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.
@@ -519,7 +533,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
@@ -527,7 +541,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:
@@ -537,21 +551,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)
@@ -563,7 +577,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()))
@@ -571,16 +585,19 @@ 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
meas_data = output.detach().cpu().numpy()
if hasattr(self, 'mask') and self.mask is not None:
mask = self.mask.detach().cpu().numpy()
else:
mask = 1
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')
@@ -599,7 +616,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]
@@ -613,8 +630,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")
@@ -625,7 +642,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':
@@ -636,5 +653,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/fancy_ptycho.py
+85 -27
View File
@@ -33,18 +33,19 @@ class FancyPtycho(CDIModel):
fourier_probe=False,
loss='amplitude mse',
units='um',
simulate_probe_translation=False
):
super(FancyPtycho, 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'])
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)
@@ -76,12 +77,11 @@ class FancyPtycho(CDIModel):
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)
shape = self.probe[0][self.detector_slice].shape
else:
background = 1e-6 * t.ones(self.probe[0].shape,
dtype=t.float32)
shape = self.probe[0].shape
background = 1e-6 * t.ones([s//oversampling for s in shape],
dtype=t.float32)
self.background = t.nn.Parameter(background)
@@ -115,6 +115,14 @@ class FancyPtycho(CDIModel):
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'):
@@ -144,7 +152,8 @@ class FancyPtycho(CDIModel):
opt_for_fft=False,
fourier_probe=False,
loss='amplitude mse',
units='um'
units='um',
simulate_probe_translation=False
):
wavelength = dataset.wavelength
@@ -248,6 +257,10 @@ class FancyPtycho(CDIModel):
# 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:
@@ -277,7 +290,8 @@ class FancyPtycho(CDIModel):
probe_support=probe_support,
fourier_probe=fourier_probe,
oversampling=oversampling,
loss=loss, units=units)
loss=loss, units=units,
simulate_probe_translation=simulate_probe_translation)
def interaction(self, index, translations, *args):
@@ -305,8 +319,15 @@ class FancyPtycho(CDIModel):
basis_prs = tools.propagators.inverse_far_field(basis_prs)
# Now we construct the probes for each shot from the basis probes
Ws = self.weights[index]
if len(self.weights[0].shape) == 0:
if 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:
@@ -317,12 +338,24 @@ class FancyPtycho(CDIModel):
# 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
@@ -352,16 +385,20 @@ class FancyPtycho(CDIModel):
self.wavelength = self.wavelength.to(*args, **kwargs)
# move the detector geometry too
det_geo = self.detector_geometry
if hasattr(det_geo, 'distance'):
if 'distance' in det_geo:
det_geo['distance'] = det_geo['distance'].to(*args, **kwargs)
if hasattr(det_geo, 'basis'):
if 'basis' in det_geo:
det_geo['basis'] = det_geo['basis'].to(*args, **kwargs)
if hasattr(det_geo, 'corner'):
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)
@@ -369,7 +406,7 @@ class FancyPtycho(CDIModel):
self.surface_normal = self.surface_normal.to(*args, **kwargs)
def sim_to_dataset(self, args_list):
def sim_to_dataset(self, args_list, calculation_width=None):
# In the future, potentially add more control
# over what metadata is saved (names, etc.)
@@ -395,9 +432,23 @@ class FancyPtycho(CDIModel):
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(indices, translations)
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,
@@ -411,11 +462,15 @@ class FancyPtycho(CDIModel):
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
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):
@@ -571,7 +626,10 @@ class FancyPtycho(CDIModel):
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()
CDTools/models/multimode_rpi.py
+359
View File
@@ -0,0 +1,359 @@
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.morphology import binary_dilation
import numpy as np
from copy import copy
__all__ = ['MultimodeRPI']
__all__ = ['RPI']
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', opt_for_fft=False, 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,
opt_for_fft=opt_for_fft,
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}
CDTools/models/polarized_fancy_ptycho.py
+260 -63
View File
@@ -3,6 +3,7 @@ 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
@@ -21,22 +22,22 @@ 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,
probe_support = None, oversampling=1,
loss='amplitude mse',units='um'):
super(FancyPtycho, self).__init__(wavelength, detector_geometry,
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 = None, translation_scale = 1, saturation=None,
probe_support = None, obj_support=None, oversampling=1,
weights = weights, translation_scale = 1, saturation=None,
probe_support = None, oversampling=1,
loss='amplitude mse',units='um')
if polarizer_offsets is None:
@@ -49,28 +50,224 @@ class PolarizedFancyPtycho(FancyPtycho):
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, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um', 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,
opt_for_fft=opt_for_fft,
loss=loss,
units=units)
# 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
#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
# WHAT IS INDEX?
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, 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
@@ -83,9 +280,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
@@ -96,17 +293,17 @@ class PolarizedFancyPtycho(FancyPtycho):
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(
prs, self.obj_support * self.obj,pix_trans,
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, :, :, :]
@@ -122,7 +319,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 +342,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 +363,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 +384,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 +402,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 +425,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 +445,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 +460,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,63 +484,63 @@ 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'),
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, 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, 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)),
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)),
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),
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)),
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)),
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)),
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()
@@ -352,7 +549,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/models/rpi.py
+45 -16
View File
@@ -4,6 +4,7 @@ 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.morphology import binary_dilation
import numpy as np
from copy import copy
@@ -49,11 +50,18 @@ class RPI(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):
obj_support=None, oversampling=1, weight_matrix=False):
super(RPI, self).__init__()
@@ -72,7 +80,7 @@ class RPI(CDIModel):
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.obj_basis = self.probe_basis * scale_factor
self.detector_slice = detector_slice
# Maybe something to include in a bit
@@ -84,10 +92,10 @@ class RPI(CDIModel):
self.mask = mask
else:
self.mask = t.tensor(mask, dtype=t.bool)
self.probe = t.tensor(probe, dtype=t.complex64)
if obj_guess.dim() == 2:
obj_guess = obj_guess[None, :, :]
@@ -96,6 +104,13 @@ class RPI(CDIModel):
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))
@@ -120,7 +135,7 @@ class RPI(CDIModel):
@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', opt_for_fft=False):
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', opt_for_fft=False, weight_matrix=False, probe_threshold=0):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -211,16 +226,24 @@ class RPI(CDIModel):
else:
raise KeyError('Initialization "' + str(initialization) + \
'" invalid - use "spectral" or "random"')
# Maybe put something here to initialize an object support based on
# a probe threshold?
obj_support=None
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)
obj_support=obj_support, oversampling=oversampling,
weight_matrix=weight_matrix)
def random_init(self, pattern):
@@ -251,8 +274,12 @@ class RPI(CDIModel):
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 = self.probe[i]
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,
@@ -275,7 +302,6 @@ class RPI(CDIModel):
output = output.unsqueeze(0).repeat(1,len(index),1,1,1)
except TypeError:
pass
return output
@@ -303,8 +329,11 @@ class RPI(CDIModel):
#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)
if self.obj.shape[0] == 1:
return factors[0] * t.sum(t.abs(self.obj[0,:,:])**2)
else:
return factors[0] * t.sum(t.abs(self.obj[0,:,:])**2) \
+ factors[1] * t.sum(t.abs(self.obj[1:,:,:])**2)
def to(self, *args, **kwargs):
super(RPI, self).to(*args, **kwargs)
@@ -336,7 +365,7 @@ class RPI(CDIModel):
plot_list = [
('Root Sum Squared Amplitude of all Probes',
lambda self, fig: p.plot_amplitude(
np.sqrt(np.sum((t.abs(self.probe)**2).cpu().numpy(),axis=0)),
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)),
('Dominant Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj[0], fig=fig,
CDTools/models/time_resolved_ptycho_calibration.py
+504
View File
@@ -0,0 +1,504 @@
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, opt_for_fft=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,
opt_for_fft=opt_for_fft,
oversampling=oversampling)
if hasattr(dataset, 'sample_info') and \
dataset.sample_info is not None and \
'orientation' in dataset.sample_info:
surface_normal = dataset.sample_info['orientation'][2]
else:
surface_normal = np.array([0., 0., 1.])
# If this information is supplied when the function is called,
# then we override the information in the .cxi file
if scattering_mode in {'t', 'transmission'}:
surface_normal = np.array([0., 0., 1.])
elif scattering_mode in {'r', 'reflection'}:
outgoing_dir = np.cross(det_basis[:, 0], det_basis[:, 1])
outgoing_dir /= np.linalg.norm(outgoing_dir)
surface_normal = outgoing_dir + np.array([0., 0., 1.])
surface_normal /= -np.linalg.norm(surface_normal)
# Next generate the object geometry from the probe geometry and
# the translations
pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations, surface_normal=surface_normal)
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
else:
background = None
# Finally, initialize the probe and object using this information
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,
}
CDTools/models/time_resolved_rpi.py
+302
View File
@@ -0,0 +1,302 @@
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.morphology import binary_dilation
import numpy as np
from copy import copy
__all__ = ['MultimodeRPI']
__all__ = ['RPI']
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', opt_for_fft=False, 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,
opt_for_fft=opt_for_fft,
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}
CDTools/tools/analysis/analysis.py
+362 -15
View File
@@ -15,7 +15,8 @@ from scipy import special
__all__ = ['orthogonalize_probes', 'standardize', 'synthesize_reconstructions',
'calc_consistency_prtf', 'calc_deconvolved_cross_correlation',
'calc_frc', 'calc_vn_entropy', 'calc_top_mode_fraction']
'calc_frc', 'calc_vn_entropy', 'calc_top_mode_fraction',
'calc_rms_error', 'calc_fidelity', 'calc_generalized_rms_error']
def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, normalize=False):
@@ -456,7 +457,7 @@ def calc_deconvolved_cross_correlation(im1, im2, im_slice=None):
return cor
def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1., limit='side'):
"""Calculates a Fourier ring correlation between two images
This function requires an input of a basis to allow for FRC calculations
@@ -479,6 +480,8 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
Number of bins to break the FRC up into
snr : float
The signal to noise ratio (for the combined information in both images) to return a threshold curve for.
limit : str
Default is 'side'. What is the highest frequency to calculate the FRC to? If 'side', it chooses the side of the Fourier transform, if 'corner' it goes fully to the corner.
Returns
-------
@@ -507,14 +510,14 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
(im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
if nbins is None:
nbins = np.max(im1[im_slice].shape) // 4
nbins = np.max(im1[im_slice].shape) // 8
f1 = t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2))
f2 = t.fft.fftshift(t.fft.fft2(im2[im_slice]),dim=(-1,-2))
cor_fft = f1 * t.conj(f2)
cor_fft = t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2)) * \
t.fft.fftshift(t.conj(t.fft.fft2(im2[im_slice])),dim=(-1,-2))
F1 = t.abs(t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2)))**2
F2 = t.abs(t.fft.fftshift(t.fft.fft2(im2[im_slice]),dim=(-1,-2)))**2
F1 = t.abs(f1)**2
F2 = t.abs(f2)**2
di = np.linalg.norm(basis[:,0])
@@ -526,21 +529,41 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
Js,Is = np.meshgrid(j_freqs,i_freqs)
Rs = np.sqrt(Is**2+Js**2)
if limit.lower().strip() == 'side':
max_i = np.max(i_freqs)
max_j = np.max(j_freqs)
frc_range = [0, max(max_i,max_j)]
elif limit.lower().strip() == 'corner':
frc_range = [0, np.max(Rs)]
else:
raise ValueError('Invalid FRC limit: choose "side" or "corner"')
numerator, bins = np.histogram(Rs, bins=nbins, range=frc_range,
weights=cor_fft.numpy())
denominator_F1, bins = np.histogram(Rs, bins=nbins, range=frc_range,
weights=F1.detach().cpu().numpy())
denominator_F2, bins = np.histogram(Rs, bins=nbins, range=frc_range,
weights=F2.detach().cpu().numpy())
n_pix, bins = np.histogram(Rs, bins=nbins, range=frc_range)
numerator, bins = np.histogram(Rs,bins=nbins,weights=cor_fft.numpy())
denominator_F1, bins = np.histogram(Rs,bins=nbins,weights=F1.detach().cpu().numpy())
denominator_F2, bins = np.histogram(Rs,bins=nbins,weights=F2.detach().cpu().numpy())
n_pix, bins = np.histogram(Rs,bins=nbins)
frc = np.abs(numerator / np.sqrt(denominator_F1*denominator_F2))
n_pix = n_pix / 4 # This is for an apodized image, apodized with a hann window
frc = np.abs(numerator) / np.sqrt(denominator_F1*denominator_F2)
# This moves from combined-image SNR to single-image SNR
snr /= 2
threshold = (snr + (2 * snr + 1) / np.sqrt(n_pix)) / \
# NOTE: I should update this to produce lots of different threshold curves
# sigma, 2sigma, 3sigma, traditional FRC 1-bit, my better one, n_pix, etc.
threshold = (snr + (2 * np.sqrt(snr) + 1) / np.sqrt(n_pix)) / \
(1 + snr + (2 * np.sqrt(snr)) / np.sqrt(n_pix))
my_threshold = np.sqrt(snr**2 + (2*snr**2 + 2*snr + 1)/n_pix) / \
np.sqrt(snr**2 + 2 * snr + 1 + 2*snr**2 / n_pix)
twosigma_threshold = 2/ np.sqrt(n_pix)
if not im_np:
bins = t.tensor(bins)
frc = t.tensor(frc)
@@ -581,6 +604,7 @@ def calc_vn_entropy(matrix):
entropy = -np.sum(special.xlogy(eig,eig))/np.sum(eig)
return entropy
def calc_top_mode_fraction(matrix):
"""Calculates the fraction of total power in the top mode of a density matrix
@@ -610,3 +634,326 @@ def calc_top_mode_fraction(matrix):
eig = np.linalg.eigh(matrix)[0]
fraction = np.max(eig) / np.sum(eig)
return fraction
def calc_rms_error(field_1, field_2, align_phases=True, normalize=False,
dims=2):
"""Calculates the root-mean-squared error between two complex wavefields
The formal definition of this function is:
output = norm * sqrt(mean(abs(field_1 - gamma * field_2)**2))
Where norm is an optional normalization factor, and gamma is an
optional phase factor which is appropriate when the wavefields suffer
from a global phase degeneracy as is often the case in diffractive
imaging.
The normalization is defined as the square root of the total intensity
contained in field_1, which is appropriate when field_1 represents a
known ground truth:
norm = sqrt(mean(abs(field_1)**2))
The phase offset is an analytic expression for the phase offset which
will minimize the RMS error between the two wavefields:
gamma = exp(1j * angle(sum(field_1 * conj(field_2))))
This implementation is stable even in cases where field_1 and field_2
are completely orthogonal.
In the definitions above, the field_n are n-dimensional wavefields. The
dimensionality of the wavefields can be altered via the dims argument,
but the default is 2 for a 2D wavefield.
Parameters
----------
field_1 : array
The first complex-valued field
field_2 : array
The second complex-valued field
align_phases : bool
Default is True, whether to account for a global phase offset
normalize : bool
Default is False, whether to normalize to the intensity of field_1
dims : (int or tuple of python:ints)
Default is 2, the number of final dimensions to reduce over.
Returns
-------
rms_error : float or t.Tensor
The RMS error, or tensor of RMS errors, depending on the dim argument
"""
sumdims = tuple(d - dims for d in range(dims))
if align_phases:
# Keepdim allows us to broadcast the result correctly when we
# multiply by the fields
gamma = t.exp(1j * t.angle(t.sum(field_1 * t.conj(field_2), dim=sumdims,
keepdim=True)))
else:
gamma = 1
if normalize:
norm = 1 / t.mean(t.abs(field_1)**2, dim=sumdims)
else:
norm = 1
difference = field_1 - gamma * field_2
return t.sqrt(norm * t.mean(t.abs(difference)**2, dim=sumdims))
def calc_fidelity(fields_1, fields_2, dims=2):
"""Calculates the fidelity between two density matrices
The fidelity is a comparison metric between two density matrices
(i.e. mutual coherence functions) that extends the idea of the
overlap to incoherent light. As a reminder, the overlap between two
fields is:
overlap = abs(sum(field_1 * field_2))**2
Whereas the fidelity is defined as:
fidelity = trace(sqrt(sqrt(dm_1) <dot> dm_2 <dot> sqrt(dm_1)))**2
where dm_n refers to the density matrix encoded by fields_n such
that dm_n = fields_n <dot> fields_<n>.conjtranspose(), sqrt
refers to the matrix square root, and <dot> is the matrix product.
This is not a practical implementation, however, as it is not feasible
to explicitly construct the matrices dm_1 and dm_2 in memory. Therefore,
we take advantage of the alternate definition based directly on the
fields_<n> parameter:
fidelity = sum(svdvals(fields_1 <dot> fields_2.conjtranspose()))**2
In the definitions above, the fields_n are regarded as collections of
wavefields, where each wavefield is by default 2-dimensional. The
dimensionality of the wavefields can be altered via the dims argument,
but the fields_n arguments must always have at least one more dimension
than the dims argument. Any additional dimensions are treated as batch
dimensions.
Parameters
----------
fields_1 : array
The first set of complex-valued field modes
fields_2 : array
The second set of complex-valued field modes
dims : int
Default is 2, the number of final dimensions to reduce over.
Returns
-------
fidelity : float or t.Tensor
The fidelity, or tensor of fidelities, depending on the dim argument
"""
fields_1 = t.as_tensor(fields_1)
fields_2 = t.as_tensor(fields_2)
mult = fields_1.unsqueeze(-dims-2) * fields_2.unsqueeze(-dims-1).conj()
sumdims = tuple(d - dims for d in range(dims))
mat = t.sum(mult,dim=sumdims)
# Because I think this is the nuclear norm squared, I would like to swap
# Out the definition for this, but I need to test it before swapping.
# It also probably makes sense to implement sqrt_fidelity separately
# because that's more important
#return t.linalg.matrix_norm(mat, ord='nuc')**2
# I think this is just the nuclear norm.
svdvals = t.linalg.svdvals(mat)
return t.sum(svdvals, dim=-1)**2
def calc_generalized_rms_error(fields_1, fields_2, normalize=False, dims=2):
"""Calculates a generalization of the root-mean-squared error between two complex wavefields
This function calculates an generalization of the RMS error which uses the
concept of fidelity to extend it to capture the error between
incoherent wavefields, defined as a mode decomposition. The extension has
several nice properties, in particular:
1) For coherent wavefields, it precisely matches the RMS error including
a correction for the global phase degeneracy (align_phases=True)
2) All mode decompositions of either field that correspond to the same
density matrix / mutual coherence function will produce the same
output
3) The error will only be zero when comparing mode decompositions that
correspond to the same density matrix.
4) Due to (2), one need not worry about the ordering of the modes,
properly orthogonalizing the modes, and it is even possible to
compare mode decompositions with different numbers of modes.
The formal definition of this function is:
output = norm * sqrt(mean(abs(fields_1)**2)
+ mean(abs(fields_2)**2)
- 2 * sqrt(fidelity(fields_1,fields_2)))
Where norm is an optional normalization factor, and the fidelity is
defined based on the mean, rather than the sum, to match the convention
for the root *mean* squared error.
The normalization is defined as the square root of the total intensity
contained in fields_1, which is appropriate when fields_1 represents a
known ground truth:
norm = sqrt(mean(abs(fields_1)**2))
In the definitions above, the fields_n are regarded as collections of
wavefields, where each wavefield is by default 2-dimensional. The
dimensionality of the wavefields can be altered via the dims argument,
but the fields_n arguments must always have at least one more dimension
than the dims argument. Any additional dimensions are treated as batch
dimensions.
Parameters
----------
fields_1 : array
The first set of complex-valued field modes
fields_2 : array
The second set of complex-valued field modes
normalize : bool
Default is False, whether to normalize to the intensity of fields_1
dims : (int or tuple of python:ints)
Default is 2, the number of final dimensions to reduce over.
Returns
-------
rms_error : float or t.Tensor
The generalized RMS error, or tensor of generalized RMS errors, depending on the dim argument
"""
fields_1 = t.as_tensor(fields_1)
fields_2 = t.as_tensor(fields_2)
npix = t.prod(t.as_tensor(fields_1.shape[-dims:],dtype=t.int32))
sumdims = tuple(d - dims - 1 for d in range(dims+1))
fields_1_intensity = t.sum(t.abs(fields_1)**2,dim=sumdims) / npix
fields_2_intensity = t.sum(t.abs(fields_2)**2,dim=sumdims) / npix
fidelity = calc_fidelity(fields_1, fields_2, dims=dims) / npix**2
result = fields_1_intensity + fields_2_intensity - 2 * t.sqrt(fidelity)
if normalize:
result /= fields_1_intensity
return t.sqrt(result)
def calc_generalized_frc(fields_1, fields_2, basis, im_slice=None, nbins=None, snr=1.):
"""Calculates a Fourier ring correlation between two images
This function requires an input of a basis to allow for FRC calculations
to be related to physical units.
Like other analysis functions, this can take input in numpy or pytorch,
and will return output in the respective format.
Parameters
----------
im1 : array
The first image, a complex or real valued array
im2 : array
The first image, a complex or real valued array
basis : array
The basis for the images, defined as is standard for datasets
im_slice : slice
Default is from 3/8 to 5/8 across the image, a slice to use in the processing.
nbins : int
Number of bins to break the FRC up into
snr : float
The signal to noise ratio (for the combined information in both images) to return a threshold curve for.
Returns
-------
freqs : array
The frequencies associated with each FRC value
FRC : array
The FRC values
threshold : array
The threshold curve for comparison
"""
im_np = False
if isinstance(fields_1, np.ndarray):
fields_1 = t.as_tensor(fields_)
im_np = True
if isinstance(fields_2, np.ndarray):
fields_2 = t.as_tensor(fields_2)
im_np = True
if isinstance(basis, np.ndarray):
basis = t.tensor(basis)
if im_slice is None:
im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5,
(im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
if nbins is None:
nbins = np.max(fields_1[...,im_slice].shape[-2:]) // 4
f1 = t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2))
f2 = t.fft.fftshift(t.fft.fft2(im2[im_slice]),dim=(-1,-2))
cor_fft = f1 * t.conj(f2)
F1 = t.abs(f1)**2
F2 = t.abs(f2)**2
di = np.linalg.norm(basis[:,0])
dj = np.linalg.norm(basis[:,1])
i_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[0],d=di))
j_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[1],d=dj))
Js,Is = np.meshgrid(j_freqs,i_freqs)
Rs = np.sqrt(Is**2+Js**2)
# This line is used to get a set of bins that matches the logic
# used by np.histogram, so that this function will match the choices
# of bin edges that comes from the non-generalized version. This also
# gets us the count on the number of pixels per bin so we can calculate
# the threshold curve
n_pix, bins = np.histogram(Rs,bins=nbins)
frc = []
for i in range(len(bins)-1):
mask = t.logical_and(Rs<bins[i+1], Rs>=bins[i])
masked_f1 = f1 * mask[...,:,:]
masked_f2 = f2 * mask[...,:,:]
numerator = t.sqrt(calc_fidelity(masked_f1, masked_f2))
denominator_f1 = t.sqrt(calc_fidelity(masked_f1, masked_f1))
denominator_f2 = t.sqrt(calc_fidelity(masked_f2, masked_f2))
frc.append(numerator / t.sqrt((denominator_f1 * denominator_f2)))
frc = np.array(frc)
# This moves from combined-image SNR to single-image SNR
snr /= 2
threshold = (snr + (2 * snr + 1) / np.sqrt(n_pix)) / \
(1 + snr + (2 * np.sqrt(snr)) / np.sqrt(n_pix))
if not im_np:
bins = t.tensor(bins)
frc = t.tensor(frc)
threshold = t.tensor(threshold)
return bins[:-1], frc, threshold
CDTools/tools/data/data.py
+9 -6
View File
@@ -93,7 +93,7 @@ def get_sample_info(cxi_file):
"""
if 'entry_1/sample_1' not in cxi_file:
return None
s1 = cxi_file['entry_1/sample_1']
metadata_attrs = ['name','description','unit_cell_group']
@@ -317,7 +317,7 @@ def get_dark(cxi_file):
return darks
def get_data(cxi_file, cut_zeroes = True):
def get_data(cxi_file, cut_zeros = True):
"""Returns an array with the full stack of detector data defined in the cxi file object
This function will make sure to check all the various places that it's
@@ -333,6 +333,8 @@ def get_data(cxi_file, cut_zeroes = True):
----------
cxi_file : h5py.File
A file object to be read
cut_zeros : bool
Default True, whether to set all negative data to zero
Returns
-------
@@ -351,9 +353,8 @@ def get_data(cxi_file, cut_zeroes = True):
raise KeyError('Data is not defined within cxi file')
data = cxi_file[pull_from][:]
# Use maximum in-place to avoid allocating any more memory than is needed
if cut_zeroes:
if cut_zeros:
np.maximum(data,0,data)
if 'axes' in cxi_file[pull_from].attrs:
@@ -642,7 +643,8 @@ def add_dark(cxi_file, dark):
d1.create_dataset('data_dark',data=dark)
def add_data(cxi_file, data, axes=None):
def add_data(cxi_file, data, axes=None, compression='gzip',
chunks=True):
"""Adds the specified data to the cxi file
It will add the data unchanged to the file, placing it in two spots:
@@ -673,7 +675,8 @@ def add_data(cxi_file, data, axes=None):
if isinstance(data, t.Tensor):
data = data.detach().cpu().numpy()
det1.create_dataset('data', data=data)
det1.create_dataset('data', data=data, compression=compression,
chunks=chunks)
data1['data'] = h5py.SoftLink('/entry_1/instrument_1/detector_1/data')
if axes is not None:
CDTools/tools/initializers/initializers.py
+1 -1
View File
@@ -633,7 +633,7 @@ def generate_subdominant_modes(dominant_mode, n_modes, circular=True):
dominant_fft = far_field(dominant_mode)
shape = dominant_mode.shape
center = ((shape[-2]-1)//2, (shape[-1]-1)//2)
center = ((shape[-2])//2, (shape[-1])//2)
i, j = np.mgrid[:shape[-2], :shape[-1]]
i = t.tensor(i - center[0]).to(dtype=dominant_fft.dtype,
CDTools/tools/interactions/interactions.py
+69 -74
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
@@ -43,7 +43,7 @@ def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.
pixel_translations : torch.Tensor
A Jx2 stack of translations in internal (i,j) pixel-space, or a single translation
"""
projection_1 = t.Tensor([[1,0,0],
[0,1,0],
[0,0,0]]).to(device=translations.device,dtype=translations.dtype)
@@ -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
@@ -423,26 +423,27 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
exit_waves : torch.Tensor
An (N)x(P)xMxL tensor of the calculated exit waves
"""
single_translation = False
if translations.dim() == 1:
translations = translations[None,:]
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)
selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2],
tr[1]:tr[1]+probe.shape[-1]]
for tr in integer_translations])
if polarized:
selections = t.stack([obj[:, :,tr[0]:tr[0]+probe.shape[-2],
tr[1]:tr[1]+probe.shape[-1]]
if not polarized:
selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2],
tr[1]:tr[1]+probe.shape[-1]]
for tr in integer_translations])
else:
selections = t.stack([obj[:, :, tr[0]:tr[0]+probe.shape[-2],
tr[1]:tr[1]+probe.shape[-1]]
for tr in integer_translations])
# Nx2x2xMxL tensor
exit_waves = []
if shift_probe:
i = t.arange(probe.shape[-2],device=probe.device,dtype=t.float32) \
@@ -454,15 +455,11 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
J = 2 * np.pi * J / probe.shape[-1]
phase_masks = t.exp(1j*(-subpixel_translations[:,0,None,None]*I
-subpixel_translations[:,1,None,None]*J))
if polarized:
phase_masks = phase_masks[..., None, :, :]
# Nx2x1xMxL tensor
# probe is (N)(P)x2xMxL tensor
fft_probe = t.fft.fftshift(t.fft.fft2(probe),dim=(-1,-2))
if multiple_modes: # Multi-mode probe
if polarized:
shifted_fft_probe = fft_probe * phase_masks[...,None,:,:,:]
@@ -470,11 +467,8 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
shifted_fft_probe = fft_probe * phase_masks[...,None,:,:]
else:
shifted_fft_probe = fft_probe * phase_masks
shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe,
dim=(-1,-2)))
if not polarized:
if multiple_modes: # Multi-mode probe
output = shifted_probe * selections[...,None,:,:]
@@ -483,12 +477,10 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
# selections: Nx2x2xMxL
# probe: Nx(P)x2x1xMxL
else:
# print('F', shift_probe)
output = polarization.apply_jones_matrix(shifted_probe, selections)
output = polarization.apply_jones_matrix(shifted_probe, selections, multiple_modes=multiple_modes)
else:
raise NotImplementedError('Object shift not yet implemented')
if single_translation:
return output[0]
else:
@@ -497,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
@@ -513,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
----------
@@ -537,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
@@ -556,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]
@@ -572,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')
@@ -587,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
@@ -601,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
@@ -618,17 +610,20 @@ 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
pad0l = (probe.shape[-2] - obj.shape[-2])//2
pad0r = probe.shape[-2] - obj.shape[-2] - pad0l
pad1l = (probe.shape[-1] - obj.shape[-1])//2
# 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
# having thought about this carefully.
pad2l = probe.shape[-2]//2 - obj.shape[-2]//2
pad2r = probe.shape[-2] - obj.shape[-2] - pad2l
pad1l = probe.shape[-1]//2 - obj.shape[-1]//2
pad1r = probe.shape[-1] - obj.shape[-1] - pad1l
fftobj = t.nn.functional.pad(fftobj, (pad1l, pad1r, pad0l, pad0r))
fftobj = t.nn.functional.pad(fftobj, (pad1l, pad1r, pad2l, pad2r))
# Again, just an inverse FFT but with an fftshift
upsampled_obj = propagators.inverse_far_field(fftobj)
CDTools/tools/plotting/__init__.py
+1
View File
@@ -1 +1,2 @@
from CDTools.tools.plotting.plotting import *
from CDTools.tools.plotting.polarized_plotting import *
CDTools/tools/plotting/plotting.py
+54 -43
View File
@@ -84,7 +84,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, interpolation=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 +94,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
----------
@@ -122,7 +122,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,
@@ -144,7 +144,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
@@ -153,9 +153,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):
@@ -183,7 +183,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:
@@ -196,14 +196,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':
@@ -219,9 +219,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
@@ -258,7 +258,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):
@@ -306,7 +306,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
@@ -418,7 +418,7 @@ def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs):
units=units, **kwargs)
def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwargs):
def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, invert_xaxis=True, **kwargs):
"""Plots a set of probe translations in a nicely formatted way
Parameters
@@ -431,6 +431,8 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa
Default is um, units to report in (assuming input in m)
lines : bool
Whether to plot lines indicating the path taken
invert_xaxis : bool
Default is True. This flips the x axis to match the convention from .cxi files of viewing the image from the beam's perspective
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
@@ -455,6 +457,9 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa
translations = translations * factor
plt.plot(translations[:,0], translations[:,1],'k.')
if invert_xaxis:
plt.gca().invert_xaxis()
if lines:
plt.plot(translations[:,0], translations[:,1],'b-', linewidth=0.5)
plt.xlabel('X (' + units + ')')
@@ -463,7 +468,7 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa
return fig
def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='probe'):
def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='probe', invert_xaxis=True):
"""Plots a set of nanomap data in a flexible way
Parameters
@@ -478,6 +483,8 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
Default is um, units to report in (assuming input in m)
convention : str
Default is 'probe', alternative is 'obj'. Whether the translations refer to the probe or object.
invert_xaxis : bool
Default is True. This flips the x axis to match the convention from .cxi files of viewing the image from the beam's perspective
Returns
-------
@@ -511,7 +518,9 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
s /= 4 # A rough value to make the size work out
plt.scatter(factor * trans[:,0],factor * trans[:,1],s=s,c=values)
if invert_xaxis:
plt.gca().invert_xaxis()
plt.gca().set_facecolor('k')
plt.xlabel('Translation x (' + units + ')')
plt.ylabel('Translation y (' + units + ')')
@@ -522,12 +531,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
@@ -560,11 +569,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
@@ -573,29 +582,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],
@@ -616,7 +625,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)
@@ -636,7 +645,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:
@@ -657,9 +666,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),
@@ -667,7 +676,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
@@ -677,7 +686,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:
@@ -688,22 +697,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")
@@ -717,7 +726,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]
@@ -735,7 +744,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':
@@ -744,7 +753,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)
@@ -755,17 +764,19 @@ 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
update(0)
return fig
CDTools/tools/plotting/polarized_plotting.py
+342
View File
@@ -0,0 +1,342 @@
import torch as t
from matplotlib import pyplot as plt
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import hsv_to_rgb
from matplotlib.widgets import Slider
from matplotlib import ticker, patheffects
from CDTools.tools.polarization import polarization
all = ['visualize_components_amplitudes',
'visualize_attenuations',
'visualize_fast_axes',
'visualize_global_axes',
'visualize_phase_ret',
'visualize_probe']
def iterator(a, b):
'''
A helper function we can use to facilitate to process of iterating over 2D arrays
'''
x, y = t.arange(a), t.arange(b)
x, y = t.meshgrid(x, y)
x, y = t.ravel(x), t.ravel(y)
return (x, y)
def plot_probe_ellipse(a=1, b=1, phase_ret=0, scale=1, x0=4, y0=5):
'''
Given a probe vector at a point (x0, y0), visualizes ellipticity
of its polarization
Parameters:
-----------
a : 1D np.array
An amplitude of the horizontal component of the probe vector
b : 1D np.array
An amplitude of the vertical component of the probe vector
phase_ret : 1D np.array
A phase difference in radians bettween the phases of the y and x components
scale : int
A scaling factor
x0, y0: 1D np.array or float
Defines the location of the vector to be plotted
Returns:
--------
x, y set of points to plot a single ellipse
'''
theta = np.linspace(0, 2*np.pi, 20)
x = x0 + scale * a * np.real(np.exp(1j * theta))
y = y0 + scale * b * np.real(np.exp(1j * (theta + phase_ret)))
return x, y
def plot_attenuations(atten_slow=1, atten_fast=1, fast_ax_angle=0, scale=1, x0=4, y0=4):
'''
Plots attenuations along the fast and slow axes
Parameters:
-----------
atten_slow: 1D np.array
Attenuation along the slow axis
atten_fast: 1D np.array
Attenuation along the fast axis
fast_ax_angle: 1D np.array
An angle between the fast horizontal and fast axes
phase_ret: 1D np.array
A difference in phases gained by the slow and the fast components
The most clockwise axis is always considered to be the fast one
scale: 1D np.array
A scaling factor
x0, y0: 1D np.array or float
Defines the location to be plotted at
Returns:
--------
x, y set of points to plot a single Jones matrix of the object
'''
angle = fast_ax_angle
theta = np.linspace(0, 2*np.pi, 20)
# collection of points to plot a fast axis
print('angle', np.rad2deg(angle))
x = x0 + scale * atten_fast * np.real(np.exp(1j * theta))
x_f = x0 + (x - x0) * np.cos(angle)
y_f = y0 + (x - x0) * np.sin(angle)
# collection of point to plot a slow axis
y = y0 + scale * atten_slow * np.real(np.exp(1j * theta))
x_s = x0 - (y - y0) * np.sin(angle)
y_s = y0 + (y - y0) * np.cos(angle)
return x_f, y_f, x_s, y_s
def plot_fast_axis(fast_ax_angle=0, scale=1, x0=4, y0=4):
'''
Plots directions of the fast axes only
'''
xf, yf, xs, ys = plot_attenuations(fast_ax_angle=fast_ax_angle, atten_fast=1, atten_slow=0, scale=scale, x0=x0, y0=y0)
return xf, yf
def plot_figures(shape, num_of_el_along_x=20, num_of_el_along_y=20,
phases=None, fast_ax_angles=None,
atten_fast=None, atten_slow=None, scale=5,
probe=False, attenuations=False, fast_axes=False):
"""
All the parameters - np.arrays of shape (shape)
"""
x_centers = np.linspace(1, shape[0] - 1, num_of_el_along_x)
y_centers = np.linspace(1, shape[1] - 1, num_of_el_along_y)
X, Y = np.meshgrid(x_centers, y_centers)
X, Y = np.ravel(X), np.ravel(Y)
xs, ys = np.array([]), np.array([])
for x, y in zip(X, Y):
k, m = int(x), int(y)
if probe:
xx, yy = plot_probe_ellipse(a=atten_fast[k, m], b=atten_slow[k, m],
phase_ret=phases[k, m], scale=scale, x0=x, y0=y)
title = 'Polarized Probe'
elif attenuations:
xf, yf, xs, ys = plot_attenuations(atten_slow=atten_slow[k, m], atten_fast=atten_fast[k, m],
fast_ax_angle=fast_ax_angles[k, m], scale=scale, x0=x, y0=y)
elif fast_axes:
xx, yy = plot_fast_axis(fast_ax_angle=fast_ax_angles[k, m], scale=scale, x0=x, y0=y)
title = 'Fast Axes'
if attenuations:
plt.plot(xf, yf, c='b')
plt.plot(xs, ys, c='b')
plt.axis('equal')
plt.title('Attenuations')
else:
plt.plot(xx, yy, c='b')
plt.axis('equal')
plt.title(title)
"""
VISUALIZATION FUNCTIONS (and helper functions)
"""
def determine_fast_axis(angle0, angle1):
# in radians
w0, w1 = t.as_tensor(angle0, dtype=t.float32), t.as_tensor(angle1, dtype=t.float32)
diff = (w0 - w1) % (2 * np.pi)
if t.tensor(0, dtype=t.float32) <= diff and diff <= t.tensor(np.pi, dtype=t.float32):
fast = 1
phase_ret = (w0 - w1) % (2 * np.pi)
gl_phase = w1
else:
fast = 0
phase_ret = (w1 - w0) % (2 * np.pi)
gl_phase = w0
if t.allclose((w0 - w1) % (2 * np.pi), phase_ret):
gl_phase = w1
fast = 1
return fast, phase_ret, gl_phase
def retrieve_obj_info(obj):
obj_t = obj.transpose(-1, -3).transpose(-2, -4)
w, v = t.linalg.eig(obj_t)
b = obj.shape[-1]
a = obj.shape[-2]
eigenvectors = t.empty(a, b, 2, dtype=t.cfloat)
ret_phases = t.empty(a, b, dtype=t.float32)
global_phases = t.empty(a, b, dtype=t.float32)
atten_fast = t.empty(a, b, dtype=t.float32)
atten_slow = t.empty(a, b, dtype=t.float32)
fast_ax_angles = t.empty(a, b, dtype=t.float32)
for k, m in zip(*iterator(a, b)):
angle0, angle1 = t.angle(w[k, m, 0]).to(dtype=t.float32), t.angle(w[k, m, 1]).to(dtype=t.float32)
fst, phase_ret, gl_phase = determine_fast_axis(angle0, angle1)
atten0, atten1 = t.abs(w[k, m, 0]), t.abs(w[k, m, 1])
ret_phases[k, m] = phase_ret
def fast(b):
if b == 0:
eigenvectors[k, m, :] = v[k, m, 0, :]
atten_fast[k, m] = atten0
atten_slow[k, m] = atten1
global_phases[k, m] = angle0
elif b == 1:
eigenvectors[k, m, :] = v[k, m, 1, :]
atten_fast[k, m] = atten1
atten_slow[k, m] = atten0
global_phases[k, m] = angle1
if t.allclose(angle0, angle1):
# then it's a linear polarizer, fast ax - the one for which attenuation is bigger
if atten0 > atten1:
fst = 0
else:
fst = 1
fast(fst)
cos = np.abs(eigenvectors[k, m, 0])
sin = np.abs(eigenvectors[k, m, 1])
if t.allclose(cos, t.zeros(1, dtype=t.float32)):
fast_ax_angles[k, m] = 90
else:
fast_ax_angles[k, m] = t.atan(sin/cos)
fast_ax_angles = np.asarray(fast_ax_angles, dtype=np.float32)
ret_phases = np.asarray(ret_phases, dtype=np.float32)
global_phases = np.asarray(global_phases, dtype=np.float32)
atten_fast = np.asarray(atten_fast, dtype=np.float32)
atten_slow = np.asarray(atten_slow, dtype=np.float32)
return fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow
def visualize_components_amplitudes(obj, rot_angle=0, logarithmic=False):
# coord_rot angle is the only angle in degrees here
def coord_rot(angle):
angle = t.as_tensor(angle, dtype=t.float32)
angle = t.deg2rad(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)
for k, m in zip(*iterator(obj.shape[-2], obj.shape[-1])):
obj[:, :, k, m] = t.matmul(coord_rot(rot_angle), obj[:, :, k, m])
components = [obj[i, j, :, :] for i, j in zip(*iterator(2, 2))]
if logarithmic:
components = [np.log(comp)/np.log(10) for comp in components]
titles = ['Amptlitudes of the a components', 'Amplitudes of the b components',
'Amplitudes of the c components', 'Amplitudes of the d components']
for i in range(4):
amplitude = np.abs(components[i])
plt.imshow(module)
plt.colorbar()
plt.title(titles[i])
def visualize_phase_ret(obj, logarithmic=False):
print('DHSBCUYLIWGBCGLWIYV')
fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow = retrieve_obj_info(obj)
if logarithmic:
ret_phases = np.log(ret_phases)/np.log(10)
plt.imshow(ret_phases)
plt.colorbar()
plt.show()
def visuallize_global_phases(obj, logarithmic=False):
fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow = retrieve_obj_info(obj)
if logarithmic:
global_phases = np.log(global_phases)/np.log(10)
plt.imshow(global_phases)
plt.colorbar()
plt.show()
def visualize_fast_axes(obj, num_of_el_along_x=20, num_of_el_along_y=20, scale=1):
fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow = retrieve_obj_info(obj)
A, B = obj.shape[-2], obj.shape[-1]
plot_figures((A, B), num_of_el_along_x=num_of_el_along_x, num_of_el_along_y=num_of_el_along_y,
fast_ax_angles=fast_ax_angles, scale=scale, fast_axes=True)
def visualize_attenuations(obj, num_of_el_along_x=20, num_of_el_along_y=20, scale=1):
fast_ax_angles, ret_phases, global_phases, atten_fast, atten_slow = retrieve_obj_info(obj)
A, B = obj.shape[-2], obj.shape[-1]
plot_figures((A, B), num_of_el_along_x=num_of_el_along_x, num_of_el_along_y=num_of_el_along_y,
phases=ret_phases, atten_slow=atten_slow, atten_fast=atten_fast,
fast_ax_angles=fast_ax_angles, scale=scale, attenuations=True)
def visualize_probe(probe, scale=1, num_of_el_along_x=20, num_of_el_along_y=20):
a = np.abs(probe[..., 0, :, :])
b = np.abs(probe[..., 1, :, :])
phases = np.angle(probe[..., 1, :, :]) - np.angle(probe[..., 0, :, :])
A, B = np.asarray(probe.shape[-2]), np.asarray(probe.shape[-1])
plot_figures((A, B), num_of_el_along_x=num_of_el_along_x, num_of_el_along_y=num_of_el_along_y,
phases=phases, atten_fast=a, atten_slow=b, scale=scale,
probe=True)
#
# def object_from_components(a, b, c, d):
# ab = t.stack((a, b), dim=-3)
# cd = t.stack((c, d), dim=-3)
# return t.stack((ab, cd), dim=-4)
#
# def object_from_quarters(a, b, c, d):
# ab = t.cat((a, b), dim=-1)
# cd = t.cat((c, d), dim=-1)
# return t.cat((ab, cd), dim=-2)
#
# def generate_birefringent_obj(shape, func_axes=None, func_ret_phases=None, func_global_phases=None, fast_axes=[0, 0, 90, 90], phases=[0, 18, 40, 18]):
# ret = [polarization.generate_birefringent_obj(fast_axis=i, phase_ret=j) for i, j in zip(fast_axes, phases)]
# A = shape[0]
# B = shape[1]
# def to_rad(angle):
# angle = t.as_tensor(angle)
# return t.deg2rad(angle)
#
# if func_axes is None:
# # 4 sets of components for each quarter
# components = [[ret[i][j][k].repeat(A//2, B//2) for j, k in zip(*iterator(2, 2))] for i in range(4)]
# # build the quarters from the components
# quarters = [object_from_components(*comp) for comp in components]
# obj = object_from_quarters(*quarters)
# else:
# X, Y = t.arange(A), t.arange(B)
# X, Y = t.meshgrid(X, Y)
# obj = t.empty(2, 2, A, B, dtype=t.cfloat)
# for i, j in zip(*iterator(A, B)):
# x, y = X[i, j], Y[i, j]
# axis = func_axes(x, y)
# ret = func_ret_phases(x, y)
# gl = func_global_phases(x, y)
# obj[:, :, i, j] = polarization.generate_birefringent_obj(fast_axis=axis, phase_ret=ret, global_phase=gl)
# w, v = t.linalg.eig(obj[:, :, i, j])
#
# # t.tensor of shape (2, 2, A, B)
# return obj
#
# def axes(x, y):
# # return (x-25)**2 + (y-25)**2
# return x * 20
#
# def phases(x, y):
# # return ((x - 5) ** 2 + (y - 5) ** 2) * 20
# return (x) * 10
#
# def glob(x, y):
# return x * 10
#
# def amp(x, y):
# return 1 + 0.05 * (x + y)
#
# def build_probe(shape, phase_ret_func=None, amps_func=None):
# probe = t.empty(2, shape[0], shape[1], dtype=t.cfloat)
# for x, y in zip(*iterator(probe.shape[-2], probe.shape[-1])):
# phase = phase_ret_func(x, y)
# phase = t.deg2rad(phase).to(dtype=t.cfloat)
# abs = amps_func(x, y).to(dtype=t.cfloat)
# probe[0, x, y] = abs
# probe[1, x, y] = abs * t.exp(phase * 1j)
#
# return probe
#
#
# obj = generate_birefringent_obj((10, 10), func_axes=axes, func_ret_phases=phases, func_global_phases=glob)
# probe = build_probe((10, 10), phase_ret_func=phases, amps_func=amp)
# # visualize_probe(probe, num_of_el_along_x=10, num_of_el_along_y=10, scale=0.1)
# visualize_fast_axes(obj, num_of_el_along_x=10, num_of_el_along_y=10, scale=0.1)
# # visualize_phase_ret(obj)
# # visuallize_global_phases(obj)
# # visualize_attenuations(obj, num_of_el_along_x=10, num_of_el_along_y=10, scale=0.3)
# plt.show()
CDTools/tools/polarization/polarization.py
+81 -108
View File
@@ -12,29 +12,32 @@ __all__ = ['apply_linear_polarizer',
'apply_half_wave_plate',
'apply_quarter_wave_plate',
'apply_circular_polarizer',
'apply_jones_matrix']
'apply_jones_matrix',
'generate_linear_polarizer',
'generate_birefringent_obj']
# Abe - split these into two functions
# Note for the future: this function should
def generate_linear_polarizer(pol_angle):
single_angle = False
pol_angle = t.as_tensor(pol_angle)
pol_angle = t.as_tensor(pol_angle).to(dtype=t.float32)
if pol_angle.dim() == 0:
pol_angle = t.unsqueeze(pol_angle,0)
single_angle = True
pol_angle_rad = t.deg2rad(pol_angle)
jones_matrices = t.stack([t.tensor([[(t.cos(p)) ** 2, t.sin(p) * t.cos(p)],
[t.sin(p) * t.cos(p), (t.sin(p)) ** 2]])
for p in polarizer])
a = t.cos(pol_angle_rad) ** 2
b = t.sin(pol_angle_rad) * t.cos(pol_angle_rad)
c = b
d = t.sin(pol_angle_rad) ** 2
ab = t.stack((a, b), dim=-1)
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):
"""
@@ -51,13 +54,17 @@ 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(polarization)
jones_matrices = generate_linear_polarizer(polarizer)
return apply_jones_matrix(probe, jones_matrices, transpose=transpose, multiple_modes=multiple_modes)
def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True):
# print('probe', probe.shape, 'jones matrix', jones_matrix.shape)
# if jones_matrix.shape == t.Size([5, 2, 2, 2, 2]):
# print('probe', probe.shape, 'jones', jones_matrix.shape)
# print('JONES', jones_matrix)
"""
Applies a given Jones matrix to the probe
@@ -66,64 +73,40 @@ 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
Assume that if the probe has a dimension (N), so does the jones matrix
"""
if multiple_modes:
if transpose:
if len(jones_matrix.shape) >= 4:
jones_matrix = jones_matrix[..., None, :, :, :, :]
else:
jones_matrix = jones_matrix[..., None, :, :, None, None]
probe = probe[..., None, :, :]
# if jones matrices do not differ from pattern to pattern
if len(probe.shape) > len(jones_matrix.shape):
jones_matrix = jones_matrix[None, ...]
jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4)
# (N)1xMxLx2x2 or (N)1x1x1x2x2
probe = probe.transpose(-1, -3).transpose(-2, -4)
output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3)
# (N)Px2xMxL
else:
if len(jones_matrix.shape) < 4:
jones_matrix = jones_matrix[..., None, :, :, :, :]
probe = t.stack((probe, probe), dim=-4)
output = t.sum(jones_matrix * probe, dim=-3)
#(N)x2xMxL
if transpose:
if jones_matrix.dim() < 4:
jones_matrix = jones_matrix[..., None, None]
if multiple_modes:
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():
jones_matrix = jones_matrix.unsqueeze(0)
# vice versa
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)
probe = probe.transpose(-1, -3).transpose(-2, -4)
output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3)
else:
if transpose:
if len(jones_matrix.shape) < 4:
jones_matrix = jones_matrix[..., None, None]
probe = probe[..., None, :, :]
# if jones matrices do not differ from pattern to pattern
if len(probe.shape) > len(jones_matrix.shape):
jones_matrix = jones_matrix[None, ...]
probe = probe.transpose(-1, -3).transpose(-2, -4)
jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4)
output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3)
raise NotImplementedError
else:
if len(jones_matrix.shape) < 4:
jones_matrix = jones_matrix[..., None, None]
probe = t.stack((probe, probe), dim=-4)
output = t.sum(jones_matrix * probe, dim=-3)
return output
def apply_phase_retardance(probe, phase_shift):
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:
----------
@@ -135,44 +118,41 @@ def apply_phase_retardance(probe, phase_shift):
Returns:
--------
probe: t.Tensor
(...)x2x1xMxL
(...)x2x1xMxL
"""
theta = t.as_tensor(phase_shift, dtype=t.float32)
theta = t.deg2rad(theta)
probe = probe.to(dtype=t.cfloat)
jones_matrix = t.tensor([[1, 0], [0, phase_shift]])
probe = probe.transpose(-1, -3).transpose(-2, -4)
polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
jones_matrix = t.tensor([[1, 0], [0, t.exp(phase_shift)]]).to(dtype=t.cfloat)
polarized = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes)
# Transpose it back
return polarized_probe.transpose(-1, -3).transpose(-2, -4)
return polarized
def apply_circular_polarizer(probe, left_polarized=True):
def apply_circular_polarizer(probe, left_polarized=True, multiple_modes=True):
"""
Applies a circular polarizer to the probe
Parameters:
----------
probe: t.Tensor
A (...)x2x1xMxL tensor representing the probe
A (...)x2xMxL tensor representing the probe
left_polarizd: bool
True for the left-polarization, False for the right
Returns:
--------
circularly polarized probe: t.Tensor
(...)x2x1xMxL
(...)x2xMxL
"""
probe = probe.to(dtype=t.cfloat)
if left_polarized:
jones_matrix = (1/2 * t.tensor([[1, -1j], [1j, 1]]))
jones_matrix = (1/2 * t.tensor([[1, -1j], [1j, 1]])).to(dtype=t.cfloat)
else:
jones_matrix = 1/2 * t.tensor([[1, 1j], [-1j, 1]])
probe = probe.transpose(-1, -3).transpose(-2, -4)
polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
jones_matrix = 1/2 * t.tensor([[1, 1j], [-1j, 1]]).to(dtype=t.cfloat)
polarized = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes)
return polarized
# Transpose it back
return polarized_probe.transpose(-1, -3).transpose(-2, -4)
def apply_quarter_wave_plate(probe, fast_axis_angle):
def apply_quarter_wave_plate(probe, fast_axis_angle, multiple_modes=True):
"""
Parameters:
----------
@@ -184,19 +164,17 @@ def apply_quarter_wave_plate(probe, fast_axis_angle):
Returns:
--------
polarized probe: t.Tensor
(...)x2x1xMxL
(...)x2x1xMxL
"""
probe = probe.to(dtype=t.cfloat)
theta = math.radians(fast_axis_angle)
exponent = t.exp(-1j * math.pi / 4 * t.ones(2, 2))
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]])
probe = probe.transpose(-1, -3).transpose(-2, -4)
polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
# Transpose it back
return polarized_probe.transpose(-1, -3).transpose(-2, -4)
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
def apply_half_wave_plate(probe, fast_axis_angle):
def apply_half_wave_plate(probe, fast_axis_angle, multiple_modes=True):
"""
Parameters:
----------
@@ -208,37 +186,32 @@ def apply_half_wave_plate(probe, fast_axis_angle):
Returns:
--------
polarized probe: t.Tensor
(...)x2x1xMxL
(...)x2x1xMxL
"""
probe = probe.to(dtype=t.cfloat)
theta = math.radians(fast_axis_angle)
exponent = t.exp(-1j * math.pi / 2 * t.ones(2, 2))
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]])
probe = probe.transpose(-1, -3).transpose(-2, -4)
polarized_probe = t.matmul(jones_matrix.to(dtype=t.cfloat), probe)
# Transpose it back
return polarized_probe.transpose(-1, -3).transpose(-2, -4)
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)
# probe = t.rand(17, 7, 2, 6, 4)
# polarizer = t.rand(7)
# out = apply_linear_polarizer(probe, polarizer)
# out2 = apply_linear_polarizer(probe, polarizer, transpose=False)
return out
# print(out.shape)
# print(out2.shape)
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
# a = t.ones(17, 8, 2, 3, 4)
# b = t.ones(2, 1, 1)
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)
# probe = t.ones(5, 2, 3, 3)
# polarizer = t.tensor([45])
# exitw = apply_linear_polarizer(probe, polarizer)
# print(exitw[:, 0, :, :])
# print('y', exitw[:, 1, :, :])
#a = t.ones(2, 4)
#print(t.sum(a, dim=1).shape)
r1 = coord_rot(-fast_axis)
r2 = coord_rot(fast_axis)
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))
CDTools/tools/propagators/propagators.py
+1
View File
@@ -382,6 +382,7 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
# Bandlimiting is not implemented in the generalized function, because it
# has a less clear meaning in that setting, so we apply it here instead
if bandlimit is not None:
# No need to multiply by 2pi
ki = 2 * np.pi * t.fft.fftfreq(shape[0],spacing[0])
kj = 2 * np.pi * t.fft.fftfreq(shape[1],spacing[1])
Ki, Kj = t.meshgrid(ki,kj)
temp_tests/simulated_dataset.py
+132
View File
@@ -0,0 +1,132 @@
import numpy as np
import torch as t
from scipy import misc
from CDTools.models import PolarizedFancyPtycho
from CDTools.datasets import PolarizedPtycho2DDataset
import CDTools
from CDTools.tools import polarization
from CDTools import tools
from matplotlib import pyplot as plt
from PIL import Image
# upolad 4 different images representing 4 components of the object
# and 2 gaaussian functionas corresponding to the probe components
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:]
print(1)
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]
translations = []
xs, ys = np.mgrid[:num_patt, :num_patt]
xs, ys = np.ravel(xs), np.ravel(ys)
for x, y in zip(xs, ys):
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, 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, 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
}
return PolarizedPtycho2DDataset(real_translations, polarizer, analyzer, patterns,
axes=("x", "y"), detector_geometry=detector_geometry, wavelength=wavelength)
# 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()
# 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)
dataset = simulate_polarized_dataset(50, 100, 10, aa, bb, cc, dd)
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()
tests/tools/test_analysis.py
+139 -1
View File
@@ -1,5 +1,6 @@
import numpy as np
from scipy import fftpack as ffts
from scipy import linalg as la
import torch as t
from itertools import combinations
@@ -239,7 +240,6 @@ def test_calc_deconvolved_cross_correlation():
assert np.allclose(test_cor_t.numpy(), np_cor)
def test_calc_frc():
obj1 = np.random.rand(270,230) + 1j * np.random.rand(270,230)
@@ -300,3 +300,141 @@ def test_calc_frc():
assert np.allclose(bins, test_bins_t.numpy())
assert np.allclose(frc, test_frc_t.numpy())
assert np.allclose(threshold, test_threshold_t.numpy())
def test_calc_rms_error():
field_1 = t.rand(14,19, dtype=t.complex64)
field_2 = t.rand(14,19, dtype=t.complex64)
# Check that the calculation is insensitive to phase
assert t.allclose(analysis.calc_rms_error(field_1, field_2),
analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2))
# And that it is sensitive to phase if we turn off the
assert not t.allclose(
analysis.calc_rms_error(field_1, field_2, align_phases=False),
analysis.calc_rms_error(field_1, np.exp(0.7j) * field_2,
align_phases=False))
# Check that the result is positive
assert analysis.calc_rms_error(field_1, field_2) > 0
# And that it is a smaller number with align_phases on
assert (analysis.calc_rms_error(field_1, field_2) <=
analysis.calc_rms_error(field_1, field_2, align_phases=False))
# Now we check against an explicit implementation:
gamma = field_1 * t.conj(field_2)
gamma /= t.abs(gamma)
# This is an alternate way of doing the calculation. Actually, would this
# be a better implementation anyway? Probably no difference tbh.
rms_error_nophase = t.sqrt((t.mean(t.abs(field_1)**2) +
t.mean(t.abs(field_2)**2) -
2 * t.abs(t.mean(field_1 * t.conj(field_2)))))
assert t.allclose(rms_error_nophase,
analysis.calc_rms_error(field_1, field_2))
rms_error_phase = t.sqrt((t.mean(t.abs(field_1)**2) +
t.mean(t.abs(field_2)**2) -
2 * t.real(t.mean(field_1 * t.conj(field_2)))))
assert t.allclose(rms_error_phase,
analysis.calc_rms_error(field_1, field_2,
align_phases=False))
# Now let's test that it works along a dimension:
field_1 = t.rand(3,14,19, dtype=t.complex64)
field_2 = t.rand(3,14,19, dtype=t.complex64)
result = analysis.calc_rms_error(field_1, field_2, normalize=True)
assert (result.shape == t.Size([3]))
for i in range(3):
assert t.allclose(analysis.calc_rms_error(field_1[i],
field_2[i],
normalize=True),
result[i])
def test_calc_fidelity():
fields_1 = t.rand(2,30,17, dtype=t.complex128)
fields_2 = t.rand(3,30,17, dtype=t.complex128)
dm_1 = t.reshape(fields_1, (2,-1))
dm_1 = t.tensordot(dm_1.transpose(0,1), dm_1.conj(), dims=1).numpy()
dm_2 = t.reshape(fields_2, (3,-1))
dm_2 = t.tensordot(dm_2.transpose(0,1), dm_2.conj(), dims=1).numpy()
inner_mat = la.sqrtm(np.dot(np.dot(la.sqrtm(dm_1),dm_2),la.sqrtm(dm_1)))
fidelity = t.as_tensor(np.abs(np.trace(inner_mat))**2)
assert t.isclose(fidelity, analysis.calc_fidelity(fields_1, fields_2))
fields_1 = t.rand(1,30,17, dtype=t.complex128)
fields_2 = t.rand(1,30,17, dtype=t.complex128)
assert t.isclose(t.abs(t.sum(fields_1*fields_2.conj()))**2,
analysis.calc_fidelity(fields_1, fields_2))
# Checking that it works with extra dimensions
fields_1 = t.rand(3,3,30,17, dtype=t.complex128)
fields_2 = t.rand(3,1,30,17, dtype=t.complex128)
field_3 = t.rand(1,30,17, dtype=t.complex128)
fidelities = analysis.calc_fidelity(fields_1, fields_2)
fidelities_2 = analysis.calc_fidelity(fields_1, field_3)
for i in range(3):
assert t.isclose(analysis.calc_fidelity(fields_1[i], fields_2[i]),
fidelities[i])
assert t.isclose(analysis.calc_fidelity(fields_1[i], field_3),
fidelities_2[i])
# Check that the diensionality argument works
fields_1 = t.rand(3,2,12, dtype=t.complex128)
fields_2 = t.rand(3,2,12, dtype=t.complex128)
assert (analysis.calc_fidelity(fields_1, fields_2, dims=1).shape
== t.Size([3]))
fields_1 = t.rand(3,2,12,4,5, dtype=t.complex128)
fields_2 = t.rand(3,2,12,4,5, dtype=t.complex128)
assert (analysis.calc_fidelity(fields_1, fields_2, dims=3).shape
== t.Size([3]))
def test_calc_generalized_rms_error():
# Test that it matches the rms error for coherent fields
fields_1 = t.rand(1,30,17, dtype=t.complex128)
fields_2 = t.rand(1,30,17, dtype=t.complex128)
assert t.isclose(analysis.calc_generalized_rms_error(fields_1, fields_2),
analysis.calc_rms_error(fields_1[0], fields_2[0],
align_phases=True))
# Test that it is independent of field order
fields_1 = t.rand(5,30,17, dtype=t.complex128)
fields_2 = t.rand(3,30,17, dtype=t.complex128)
fields_3 = fields_2.flip(0)
assert t.isclose(analysis.calc_generalized_rms_error(fields_1, fields_2),
analysis.calc_generalized_rms_error(fields_1, fields_3))
# Test with leading dimensions
fields_1 = t.rand(3,4,2,10,17, dtype=t.complex128)
fields_2 = t.rand(3,4,3,10,17, dtype=t.complex128)
assert (analysis.calc_generalized_rms_error(fields_1, fields_2).shape
== t.Size([3,4]))
# And test with different number of dimensions dims
# Test that it is independent of field order
fields_1 = t.rand(3,6,17, dtype=t.complex128)
fields_2 = t.rand(3,1,17, dtype=t.complex128)
fields_3 = fields_2.flip(0)
assert (analysis.calc_generalized_rms_error(fields_1, fields_2, dims=1).shape == t.Size([3]))
tests/tools/test_polarization.py
+524 -341
View File
@@ -1,15 +1,17 @@
import numpy as np
import torch as t
from CDTools.tools.polarization import apply_linear_polarizer, generate_linear_polarizer
from CDTools.tools.polarization import apply_jones_matrix as jones
# Abe - I removed all the imports that didn't need to be here.
# Abe - A few issues. First, you could just write "from math import cos, sin"
# Second, we already have numpy imported, so better to use np.cos and np.sin
from math import cos as cos, sin as sin
# from math import cos as cos, sin as sin
# numpy also has np.pi and np.deg2rad
import math
from numpy import cos, sin, deg2rad
# Further comments:
#
@@ -34,408 +36,589 @@ angle = 87
angle_2 = angle - 45
def polarizer(angle):
theta = math.radians(angle)
polarizer = t.tensor([[(cos(theta)) ** 2, sin(2 * theta) / 2], [sin(2 * theta) / 2, sin(theta) ** 2]]).to(dtype=t.cfloat)
return polarizer
theta = deg2rad(angle)
polarizer = t.tensor([[(cos(theta)) ** 2, sin(2 * theta) / 2], [sin(2 * theta) / 2, sin(theta) ** 2]]).to(dtype=t.cfloat)
return polarizer
exponent = t.exp(-1j * math.pi / 4 * t.ones(2, 2))
theta2 = math.radians(angle_2)
exponent = t.exp(-1j * np.pi / 4 * t.ones(2, 2)).to(dtype=t.cfloat)
theta2 = deg2rad(angle_2)
quarter_plate = t.tensor([[(cos(theta2))**2 + 1j * (sin(theta2))**2, (1 - 1j) * sin(theta2) * cos(theta2)],
[(1 - 1j) * sin(theta2) * cos(theta2), (sin(theta2))**2 + 1j * (cos(theta2))**2]])
[(1 - 1j) * sin(theta2) * cos(theta2), (sin(theta2))**2 + 1j * (cos(theta2))**2]]).to(dtype=t.cfloat)
def build_from_quarters(jones1, jones2, jones3, jones4):
x = t.cat((t.stack((jones1, jones1), dim=-1), t.stack((jones2, jones2), dim=-1)), dim=-1)
y = t.cat((t.stack((jones3, jones3), dim=-1), t.stack((jones4, jones4), dim=-1)), dim=-1)
x = t.stack((x, x), dim=-2)
y = t.stack((y, y), dim=-2)
return t.cat((x, y), dim=-2).to(dtype=t.cfloat)
x = t.cat((t.stack((jones1, jones1), dim=-1), t.stack((jones2, jones2), dim=-1)), dim=-1)
y = t.cat((t.stack((jones3, jones3), dim=-1), t.stack((jones4, jones4), dim=-1)), dim=-1)
x = t.stack((x, x), dim=-2)
y = t.stack((y, y), dim=-2)
return t.cat((x, y), dim=-2).to(dtype=t.cfloat)
jones_plate = t.matmul(quarter_plate, polarizer(angle))
jones0 = polarizer(0)
jones90 = polarizer(90)
jones45 = polarizer(45)
'''
after applying the polarizer and the quarter_plate, the probe should get circularly polarized
probe: no multiple modes, 1 diffr pattern
2xMxL
jones_matrix: same jones matrix applied to all the pixels
2x2
'''
transpose = True
def test_apply_jones_matrix_no_modes_no_mult_patterns_one_jones_matr():
probe = t.rand(2, 3, 4, dtype=t.cfloat)
print(polarizer)
out = jones(jones(probe, polarizer(angle), multiple_modes=False, transpose=transpose),
quarter_plate, multiple_modes=False, transpose=transpose)
print('expected shape:(2, 3, 4)')
print('actual:', out.shape)
print('simulated:', out)
assert out.shape == t.Size((2,3,4))
def test_apply_jones_matrix_no_modes_no_mult_patterns_one_jones_matr():
'''
after applying the polarizer and the quarter_plate, the probe should get circularly polarized
probe: no multiple modes, 1 diffr pattern
2xMxL
jones_matrix: same jones matrix applied to all the pixels
2x2
'''
probe = t.rand(2, 3, 4, dtype=t.cfloat)
print(polarizer)
out = jones(jones(probe, polarizer(angle), multiple_modes=False, transpose=transpose),
quarter_plate, multiple_modes=False, transpose=transpose)
print('expected shape:(2, 3, 4)')
print('actual:', out.shape)
print('simulated:', out)
assert np.allclose(np.real(out[0]), np.imag(out[1]))
assert np.allclose(np.real(out[0]), np.imag(out[1]))
assert out.shape == t.Size((2, 3, 4))
'''probe: no multiple modes, 1 diffr pattern
2xMxL = 2x4x4
jones_matrix: jones matrices differ from pixel to pixel
2x2xMxL = 2x2x4x4
4 quarters:
1: [:, :, :-2, :-2] - circular_polarizer,
2: [:, :, :-2, -2:] - 0
3: [:, :, -2:, :-2] - 90
4: [:, :, -2:, -2:] - 45
'''
def test_generate_linear_polarizer():
pol_angles = [0, 45, 90]
pol_angle1 = 45
pol_angle2 = t.tensor(90)
pol_angle3 = t.tensor([0])
pols = generate_linear_polarizer(pol_angles)
pol1 = generate_linear_polarizer(pol_angle1)
pol2 = generate_linear_polarizer(pol_angle2)
pol3 = generate_linear_polarizer(pol_angle3)
print('polarizers 0, 45, 90 (1D tensor) shape:', pols.shape)
print('shape of the polarizer generated from int:', pol1.shape)
print('shape of the polarizer generated from 0D tensor:', pol2.shape)
print('shape of the linear polarizer generated from t.Size(0) tensor:', pol3.shape)
print('90', pol2)
print(jones90)
probe = t.ones(4, 4)
jones_m = [jones0, jones45, jones90]
jones_m = t.stack([matr for matr in jones_m])
assert pols.shape == t.Size((3, 2, 2))
assert pol1.shape == t.Size((2, 2))
assert pol2.shape == t.Size((2, 2))
assert pol3.shape == t.Size((1, 2, 2))
assert t.allclose(pol1, jones45)
def test_apply_jones_matrix_no_modes_no_mult_patterns_diff_jones_matr():
jones_matr = build_from_quarters(jones_plate, jones0, jones90, jones45)
probe = t.ones(2, 4, 4).to(dtype=t.cfloat)
print('jones:', jones_matr)
# print('jones:', jones_matr)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
# jones -> (2,2,x,y), probe, output probe
# interaction ->
print('expected shape:(2, 4, 4)')
print('simulated:', out.shape)
print('simulated:', out)
def test_apply_jones_matrix_no_modes_no_mult_patterns_diff_jones_matr():
'''probe: no multiple modes, 1 diffr pattern
2xMxL = 2x4x4
jones_matrix: jones matrices differ from pixel to pixel
2x2xMxL = 2x2x4x4
4 quarters:
1: [:, :, :-2, :-2] - circular_polarizer,
2: [:, :, :-2, -2:] - 0
3: [:, :, -2:, :-2] - 90
4: [:, :, -2:, -2:] - 45
'''
jones_matr = build_from_quarters(jones_plate, jones0, jones90, jones45)
probe = t.ones(2, 4, 4).to(dtype=t.cfloat)
print('jones:', jones_matr)
# print('jones:', jones_matr)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
# Example of using multiple asserts
assert np.allclose(np.real(out[0, :-2, :-2]), np.imag(out[1, :-2, :-2]))
assert t.allclose(out[0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, -2:, -2:], out[0, -2:, -2:])
# jones -> (2,2,x,y), probe, output probe
# interaction ->
print('expected shape:(2, 4, 4)')
print('simulated:', out.shape)
print('simulated:', out)
# Example of using multiple asserts
assert np.allclose(np.real(out[0, :-2, :-2]), np.imag(out[1, :-2, :-2]))
assert t.allclose(out[0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, -2:, -2:], out[0, -2:, -2:])
assert out.shape == t.Size((2, 4, 4))
'''
probe: no multiple modes, multiple diffr patterns
Nx2xMxL = 3x2x4x4
jones_matrix: same jones matrix applied to all the pixels
Nx2x2 = 3x2x2
3 different matrices for each probe:
1: 0
2: 45
3: 90
'''
def test_apply_jones_matrix_no_modes_mult_patterns_one_jones_matr():
probe = t.ones(2, 4, 4, dtype=t.cfloat)
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
jones_matr = t.stack(([polarizer(angle) for angle in [0, 45, 90]]), dim=0)
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
'''
probe: no multiple modes, multiple diffr patterns
Nx2xMxL = 3x2x4x4
jones_matrix: same jones matrix applied to all the pixels
Nx2x2 = 3x2x2
3 different matrices for each probe:
1: 0
2: 45
3: 90
'''
probe = t.ones(2, 4, 4, dtype=t.cfloat)
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
jones_matr = t.stack(([polarizer(angle) for angle in [0, 45, 90]]), dim=0)
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
print('expected shape: (3, 2, 4, 4)')
print('actual:', out.shape)
print('simulated:', out)
print('expected shape: (3, 2, 4, 4)')
print('actual:', out.shape)
print('simulated:', out)
assert (t.allclose(out[0, 0, :, :], t.ones(4, 4, dtype=t.cfloat))
and t.allclose(out[0, 1, :, :], t.zeros(4, 4, dtype=t.cfloat))
and t.allclose(out[1, 0, :, :], out[1, 1])
and t.allclose(out[2, 0, :, :], 3* t.zeros(4, 4, dtype=t.cfloat))
and t.allclose(out[2, 1, :, :], 3 * t.ones(4, 4, dtype=t.cfloat)))
assert t.allclose(out[0, 0, :, :], t.ones(4, 4, dtype=t.cfloat))
assert t.allclose(out[0, 1, :, :], t.zeros(4, 4, dtype=t.cfloat))
assert t.allclose(out[1, 0, :, :], out[1, 1])
assert t.allclose(out[2, 0, :, :], 3* t.zeros(4, 4, dtype=t.cfloat))
assert t.allclose(out[2, 1, :, :], 3 * t.ones(4, 4, dtype=t.cfloat))
assert out.shape == t.Size((3, 2, 4, 4))
'''
probe: no multiple modes, multiple diffr pattern
Nx2xMxL = 3x2x4x4
jones_matrix: jones matrices differ from pixel to pixel
Nx2x2xMxL = 3x2x2x4x4
jones matrix for the 1st pattern:
1: quat plate, 2: 90, 3: 0, 4: 45
jones matrix for the 2nd pattern:
1: 0, 2: 45, 3: 90, 4: quat plate
jones matrix for the 3rd pattern:
1: 90, 2: 0, 3: 45, 4: quat plate
'''
def test_apply_jones_matrix_no_modes_mult_patterns_diff_jones_matr():
jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45),
build_from_quarters(jones0, jones45, jones90, jones_plate),
build_from_quarters(jones90, jones0, jones45, jones_plate)]
jones_matr = t.stack(([i for i in jones_m]), dim=0)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
'''
probe: no multiple modes, multiple diffr pattern
Nx2xMxL = 3x2x4x4
jones_matrix: jones matrices differ from pixel to pixel
Nx2x2xMxL = 3x2x2x4x4
jones matrix for the 1st pattern:
1: quat plate, 2: 90, 3: 0, 4: 45
jones matrix for the 2nd pattern:
1: 0, 2: 45, 3: 90, 4: quat plate
jones matrix for the 3rd pattern:
1: 90, 2: 0, 3: 45, 4: quat plate
'''
jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45),
build_from_quarters(jones0, jones45, jones90, jones_plate),
build_from_quarters(jones90, jones0, jones45, jones_plate)]
jones_matr = t.stack(([i for i in jones_m]), dim=0)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
print('expected shape: (3, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated:', out)
print('expected shape: (3, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated:', out)
assert (np.allclose(np.real(out[0, 0, :-2, :-2]), np.imag(out[0, 1, :-2, :-2]))
and t.allclose(out[0, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[0, 1, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[0, 0, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[0, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:])
assert np.allclose(np.real(out[0, 0, :-2, :-2]), np.imag(out[0, 1, :-2, :-2]))
assert t.allclose(out[0, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 1, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 0, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:])
and t.allclose(out[1, 0, :-2, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[1, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[1, 0, :-2, -2:], out[1, 1, :-2, -2:])
and t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[1, 1, -2:, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
and np.allclose(np.real(out[1, 0, -2:, -2:]), np.real(out[1, 0, -2:, -2:]))
assert t.allclose(out[1, 0, :-2, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 0, :-2, -2:], out[1, 1, :-2, -2:])
assert t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 1, -2:, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
assert np.allclose(np.real(out[1, 0, -2:, -2:]), np.real(out[1, 0, -2:, -2:]))
and t.allclose(out[2, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[2, 1, :-2, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[2, 0, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[2, 0, -2:, :-2], out[2, 1, -2:, :-2])
and np.allclose(np.real(out[2, 0, -2:, -2:]), np.real(out[2, 0, -2:, -2:])))
assert t.allclose(out[2, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 1, :-2, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 0, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 0, -2:, :-2], out[2, 1, -2:, :-2])
assert np.allclose(np.real(out[2, 0, -2:, -2:]), np.real(out[2, 0, -2:, -2:]))
assert out.shape == t.Size((3, 2, 4, 4))
'''
probe: multiple modes, 1 diffr pattern
Px2xMxL = 2x2x3x4
jones_matrix: same jones matrix applied to all the pixels
2x2 - quarter plate
'''
def test_apply_jones_matrix_mult_modes_1_pattern_one_jones_matr():
probe = t.rand(2, 2, 3, 4, dtype=t.cfloat)
out = jones(jones(probe, polarizer(angle), multiple_modes=True, transpose=transpose), quarter_plate, multiple_modes=True, transpose=transpose)
'''
probe: multiple modes, 1 diffr pattern
Px2xMxL = 2x2x3x4
jones_matrix: same jones matrix applied to all the pixels
2x2 - quarter plate
'''
probe = t.rand(2, 2, 3, 4, dtype=t.cfloat)
out = jones(jones(probe, polarizer(angle), multiple_modes=True, transpose=transpose), quarter_plate, multiple_modes=True, transpose=transpose)
print('expected shape: (2, 2, 3, 4)')
print('actual:', out.shape)
print('simulated:', out)
print('expected shape: (2, 2, 3, 4)')
print('actual:', out.shape)
print('simulated:', out)
assert (np.allclose(np.real(out[0, 0, :, :]), np.imag(out[0, 1, :, :]))
and np.allclose(np.real(out[1, 0, :, :]), np.imag(out[1, 1, :, :])))
assert np.allclose(np.real(out[0, 0, :, :]), np.imag(out[0, 1, :, :]))
assert np.allclose(np.real(out[1, 0, :, :]), np.imag(out[1, 1, :, :]))
assert out.shape == t.Size((2, 2, 3, 4))
'''
probe: multiple modes, 1 diffr pattern
Px2xMxL = 3x2x4x4
jones_matrix: jones matrices differ from pixel to pixel
2xMxL = 2x4x4
4 quarters:
1: [:, :, :-2, :-2] - circular_polarizer,
2: [:, :, :-2, -2:] - 0
3: [:, :, -2:, :-2] - 90
4: [:, :, -2:, -2:] - 45
'''
def test_apply_jones_matrix_mult_modes_1_pattern_diff_jones_matr():
jones_matr = build_from_quarters(jones_plate, jones0, jones90, jones45)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
'''
probe: multiple modes, 1 diffr pattern
Px2xMxL = 3x2x4x4
jones_matrix: jones matrices differ from pixel to pixel
2xMxL = 2x4x4
4 quarters:
1: [:, :, :-2, :-2] - circular_polarizer,
2: [:, :, :-2, -2:] - 0
3: [:, :, -2:, :-2] - 90
4: [:, :, -2:, -2:] - 45
'''
jones_matr = build_from_quarters(jones_plate, jones0, jones90, jones45)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
print('expected shape: (3, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated:', out)
print('expected shape: (3, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated:', out)
assert (np.allclose(np.real(out[0, 0, :-2, :-2]), np.imag(out[0, 1, :-2, :-2]))
and t.allclose(out[0, 0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[0, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[0, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[0, 1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:])
assert np.allclose(np.real(out[0, 0, :-2, :-2]), np.imag(out[0, 1, :-2, :-2]))
assert t.allclose(out[0, 0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:])
and np.allclose(np.real(out[1, 0, :-2, :-2]), np.imag(out[1, 1, :-2, :-2]))
and t.allclose(out[1, 0, :-2, -2:], 2 * t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[1, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[1, 1, -2:, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[1, 0, -2:, -2:], out[1, 1, -2:, -2:])
assert np.allclose(np.real(out[1, 0, :-2, :-2]), np.imag(out[1, 1, :-2, :-2]))
assert t.allclose(out[1, 0, :-2, -2:], 2 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 1, -2:, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 0, -2:, -2:], out[1, 1, -2:, -2:])
and np.allclose(np.real(out[2, 0, :-2, :-2]), np.imag(out[2, 1, :-2, :-2]))
and t.allclose(out[2, 0, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[2, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[2, 1, -2:, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[2, 0, -2:, -2:], out[2, 1, -2:, -2:]))
assert np.allclose(np.real(out[2, 0, :-2, :-2]), np.imag(out[2, 1, :-2, :-2]))
assert t.allclose(out[2, 0, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 1, -2:, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 0, -2:, -2:], out[2, 1, -2:, -2:])
assert out.shape == t.Size((3, 2, 4, 4))
'''
probe: multiple modes, multiple diffr patterns
NxPx2xMxL = 3x7x2x3x4
jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern)
Nx2x2 = 3x2x2
3 different matrices for each probe in one mode:
1: 0
2: 45
3: 90
'''
def test_apply_jones_matrix_mult_modes_mult_pattern_one_jones_matr():
probe = t.ones(2, 3, 4, dtype=t.cfloat)
# 1st mode
probe_mode1 = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
# 2nd mode
probe_mode2 = 10 * t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
probe = t.stack(([probe_mode1 * 10 ** i for i in range(7)]), dim=1)
jones_matr = t.stack(([polarizer(angle) for angle in [0, 45, 90]]), dim=0)
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
'''
probe: multiple modes, multiple diffr patterns
NxPx2xMxL = 3x7x2x3x4
jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern)
Nx2x2 = 3x2x2
3 different matrices for each probe in one mode:
1: 0
2: 45
3: 90
'''
probe = t.ones(2, 3, 4, dtype=t.cfloat)
# 1st mode
probe_mode1 = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
# 2nd mode
probe_mode2 = 10 * t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
probe = t.stack(([probe_mode1 * 10 ** i for i in range(7)]), dim=1)
jones_matr = t.stack(([polarizer(angle) for angle in [0, 45, 90]]), dim=0)
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
print('expected shape: (3, 7, 2, 3, 4)')
print('actual:', out.shape)
print('simulated (patterns in one mode):', out[:, 6, :, :, :])
# we'll be checking only one mode
assert (t.allclose(out[0, 6, 0, :, :], (10**6) * t.ones(3, 4, dtype=t.cfloat))
and t.allclose(out[0, 6, 1, :, :], t.zeros(3, 4, dtype=t.cfloat))
and t.allclose(out[1, 6, 0, :, :], out[1, 6, 1, :, :])
and t.allclose(out[2, 6, 0, :, :], 3 * (10**6) * t.zeros(3, 4, dtype=t.cfloat))
and t.allclose(out[2, 6, 1, :, :], 3 * (10**6) * t.ones(3, 4, dtype=t.cfloat)))
print('expected shape: (3, 7, 2, 3, 4)')
print('actual:', out.shape)
print('simulated (patterns in one mode):', out[:, 6, :, :, :])
# we'll be checking only one mode
assert t.allclose(out[0, 6, 0, :, :], (10**6) * t.ones(3, 4, dtype=t.cfloat))
assert t.allclose(out[0, 6, 1, :, :], t.zeros(3, 4, dtype=t.cfloat))
assert t.allclose(out[1, 6, 0, :, :], out[1, 6, 1, :, :])
assert t.allclose(out[2, 6, 0, :, :], 3 * (10**6) * t.zeros(3, 4, dtype=t.cfloat))
assert t.allclose(out[2, 6, 1, :, :], 3 * (10**6) * t.ones(3, 4, dtype=t.cfloat))
assert out.shape == t.Size((3, 7, 2, 3, 4))
'''
probe: multiple modes, multiple diffr pattern
NxPx2xMxL = 3x4x2x4x4
jones_matrix: jones matrices differ from pixel to pixel
Nx2x2xMxL = 3x2x2x4x4
(differs across the patterns in each mode)
jones matrix for the 1st pattern:
1: quat plate, 2: 90, 3: 0, 4: 45
jones matrix for the 2nd pattern:
1: 0, 2: 45, 3: 90, 4: quat plate
jones matrix for the 3rd pattern:
1: 90, 2: 0, 3: 45, 4: quat plate
'''
def test_apply_jones_matrix_mult_modes_mult_patterns_diff_jones_matr():
jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45),
build_from_quarters(jones0, jones45, jones90, jones_plate),
build_from_quarters(jones90, jones0, jones45, jones_plate)]
jones_matr = t.stack(([i for i in jones_m]), dim=0)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
# one mode
probe_mode = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
probe = t.stack(([probe_mode * (10 ** i) for i in range(4)]), dim=1)
print('probe:', probe.shape)
print('jones:', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
'''
probe: multiple modes, multiple diffr pattern
NxPx2xMxL = 3x4x2x4x4
jones_matrix: jones matrices differ from pixel to pixel
Nx2x2xMxL = 3x2x2x4x4
(differs across the patterns in each mode)
jones matrix for the 1st pattern:
1: quat plate, 2: 90, 3: 0, 4: 45
jones matrix for the 2nd pattern:
1: 0, 2: 45, 3: 90, 4: quat plate
jones matrix for the 3rd pattern:
1: 90, 2: 0, 3: 45, 4: quat plate
'''
jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45),
build_from_quarters(jones0, jones45, jones90, jones_plate),
build_from_quarters(jones90, jones0, jones45, jones_plate)]
jones_matr = t.stack(([i for i in jones_m]), dim=0)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
# one mode
probe_mode = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
probe = t.stack(([probe_mode * (10 ** i) for i in range(4)]), dim=1)
print('probe:', probe.shape)
print('jones:', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
print('expected shape: (3, 2, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated patterns in one mode:', out[:, 3, :, :, :])
print('expected shape: (3, 4, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated patterns in one mode:', out[:, 3, :, :, :])
o = 10 ** 3
# we'll be checking only one mode (4th)
assert (np.allclose(np.real(out[0, 3, 0, :-2, :-2]), np.imag(out[0, 3, 1, :-2, :-2]))
and t.allclose(out[0, 3, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[0, 3, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[0, 3, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[0, 3, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[0, 3, 0, -2:, -2:], out[0, 3, 1, -2:, -2:])
o = 10 ** 3
# we'll be checking only one mode (4th)
assert np.allclose(np.real(out[0, 3, 0, :-2, :-2]), np.imag(out[0, 3, 1, :-2, :-2]))
assert t.allclose(out[0, 3, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 3, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 3, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 3, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 3, 0, -2:, -2:], out[0, 3, 1, -2:, -2:])
and t.allclose(out[1, 3, 0, :-2, :-2], 2 * o * t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[1, 3, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[1, 3, 0, :-2, -2:], out[1, 3, 1, :-2, -2:])
and t.allclose(out[1, 3, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[1, 3, 1, -2:, :-2], 2 * o * t.ones(2, 2, dtype=t.cfloat))
and np.allclose(np.real(out[1, 3, 0, -2:, -2:]), np.real(out[1, 3, 0, -2:, -2:]))
assert t.allclose(out[1, 3, 0, :-2, :-2], 2 * o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 3, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 3, 0, :-2, -2:], out[1, 3, 1, :-2, -2:])
assert t.allclose(out[1, 3, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 3, 1, -2:, :-2], 2 * o * t.ones(2, 2, dtype=t.cfloat))
assert np.allclose(np.real(out[1, 3, 0, -2:, -2:]), np.real(out[1, 3, 0, -2:, -2:]))
and t.allclose(out[2, 3, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[2, 3, 1, :-2, :-2], 3 * o * t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[2, 3, 0, :-2, -2:], 3 * o * t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[2, 3, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[2, 3, 0, -2:, :-2], out[2, 3, 1, -2:, :-2])
and np.allclose(np.real(out[2, 3, 0, -2:, -2:]), np.real(out[2, 3, 0, -2:, -2:])))
assert t.allclose(out[2, 3, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 3, 1, :-2, :-2], 3 * o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 3, 0, :-2, -2:], 3 * o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 3, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 3, 0, -2:, :-2], out[2, 3, 1, -2:, :-2])
assert np.allclose(np.real(out[2, 3, 0, -2:, -2:]), np.real(out[2, 3, 0, -2:, -2:]))
assert out.shape == t.Size((3, 4, 2, 4, 4))
'''
probe: no multiple modes, multiple diffr patterns
Nx2xMxL = 3x2x4x4
jones_matrix: same jones matrix applied to all the pixels
2x2 = 2x2 - a quarter waveplate
'''
def test_apply_jones_matrix_no_modes_mult_patterns_one_jones_matr_1():
probe = t.ones(2, 4, 4, dtype=t.cfloat)
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
jones_matr = jones_plate
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
'''
probe: no multiple modes, multiple diffr patterns
Nx2xMxL = 3x2x4x4
jones_matrix: same jones matrix applied to all the pixels
2x2 = 2x2 - a quarter waveplate
'''
probe = t.ones(2, 4, 4, dtype=t.cfloat)
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
jones_matr = jones_plate
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
print('expected shape: (3, 2, 4, 4)')
print('actual:', out.shape)
print('simulated:', out)
print('expected shape: (3, 2, 4, 4)')
print('actual:', out.shape)
print('simulated:', out)
assert np.allclose(np.real(out[:, 0, :, :]), np.imag(out[:, 1, :, :]))
assert np.allclose(np.real(out[:, 0, :, :]), np.imag(out[:, 1, :, :]))
assert out.shape == t.Size((3, 2, 4, 4))
'''
probe: no multiple modes, multiple diffr pattern
Nx2xMxL = 3x2x4x4
jones_matrix: jones matrices differ from pixel to pixel
2x2xMxL = 2x2x4x4
1: quat plate, 2: 90, 3: 0, 4: 45
'''
def test_apply_jones_matrix_no_modes_mult_patterns_diff_jones_matr_1():
jones_matr = build_from_quarters(jones_plate, jones90, jones0, jones45)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
'''
probe: no multiple modes, multiple diffr pattern
Nx2xMxL = 3x2x4x4
jones_matrix: jones matrices differ from pixel to pixel
2x2xMxL = 2x2x4x4
1: quat plate, 2: 90, 3: 0, 4: 45
print('expected shape: (3, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated:', out)
'''
jones_matr = build_from_quarters(jones_plate, jones90, jones0, jones45)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
print('expected shape: (3, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated:', out)
# check the 3rd pattern
assert (np.allclose(np.real(out[2, 0, :-2, :-2]), np.imag(out[2, 1, :-2, :-2]))
and t.allclose(out[2, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[2, 1, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[2, 0, -2:, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[2, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[2, 0, -2:, -2:], out[2, 1, -2:, -2:]))
assert np.allclose(np.real(out[2, 0, :-2, :-2]), np.imag(out[2, 1, :-2, :-2]))
assert t.allclose(out[2, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 1, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 0, -2:, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 0, -2:, -2:], out[2, 1, -2:, -2:])
assert out.shape == t.Size((3, 2, 4, 4))
'''
probe: multiple modes, multiple diffr patterns
NxPx2xMxL = 3x7x2x3x4
jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern)
2x2 = 2x2
quarter wave plate
'''
def test_apply_jones_matrix_mult_modes_mult_pattern_one_jones_matr_1():
probe = t.ones(2, 3, 4, dtype=t.cfloat)
# 1st mode
probe_mode1 = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
probe = t.stack(([probe_mode1 * 10 ** i for i in range(7)]), dim=1)
jones_matr = jones_plate
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
'''
probe: multiple modes, multiple diffr patterns
NxPx2xMxL = 3x7x2x3x4
jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern)
2x2 = 2x2
quarter wave plate
'''
probe = t.ones(2, 3, 4, dtype=t.cfloat)
# 1st mode
probe_mode1 = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
probe = t.stack(([probe_mode1 * 10 ** i for i in range(7)]), dim=1)
jones_matr = jones_plate
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
print('expected shape: (3, 7, 2, 3, 4)')
print('actual:', out.shape)
print('simulated (patterns in one mode):', out[:, 6, :, :, :])
# we'll be checking only one mode
assert np.allclose(np.real(out[:, :, 0, :, :]), np.imag(out[:, :, 1, :, :]))
print('expected shape: (3, 7, 2, 3, 4)')
print('actual:', out.shape)
print('simulated (patterns in one mode):', out[:, 6, :, :, :])
# we'll be checking only one mode
assert np.allclose(np.real(out[:, :, 0, :, :]), np.imag(out[:, :, 1, :, :]))
assert out.shape == t.Size((3, 7, 2, 3, 4))
'''
probe: multiple modes, multiple diffr pattern
NxPx2xMxL = 3x4x2x4x4
jones_matrix: jones matrices differ from pixel to pixel
2x2xMxL = 2x2x4x4
jones matrix:
1: quat plate, 2: 90, 3: 0, 4: 45
'''
def test_apply_jones_matrix_mult_modes_mult_patterns_diff_jones_matr_1():
jones_matr = build_from_quarters(jones_plate, jones90, jones0, jones45)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
# one mode
probe_mode = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
probe = t.stack(([probe_mode * (10 ** i) for i in range(4)]), dim=1)
print('probe:', probe.shape)
print('jones:', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
'''
probe: multiple modes, multiple diffr pattern
NxPx2xMxL = 3x4x2x4x4
jones_matrix: jones matrices differ from pixel to pixel
2x2xMxL = 2x2x4x4
jones matrix:
1: quat plate, 2: 90, 3: 0, 4: 45
'''
jones_matr = build_from_quarters(jones_plate, jones90, jones0, jones45)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
# one mode
probe_mode = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
probe = t.stack(([probe_mode * (10 ** i) for i in range(4)]), dim=1)
print('probe:', probe.shape)
print('jones:', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
print('expected shape: (3, 2, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated patterns in one mode:', out[:, 3, :, :, :])
print('expected shape: (3, 4, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated patterns in one mode:', out[:, 3, :, :, :])
o = 10 ** 2
# we'll be checking only one mode (3th)
assert (np.allclose(np.real(out[0, 2, 0, :-2, :-2]), np.imag(out[0, 2, 1, :-2, :-2]))
and t.allclose(out[0, 2, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[0, 2, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[0, 2, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
and t.allclose(out[0, 2, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
and t.allclose(out[0, 2, 0, -2:, -2:], out[0, 2, 1, -2:, -2:]))
o = 10 ** 2
# we'll be checking only one mode (3th)
assert np.allclose(np.real(out[0, 2, 0, :-2, :-2]), np.imag(out[0, 2, 1, :-2, :-2]))
assert t.allclose(out[0, 2, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 2, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 2, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 2, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 2, 0, -2:, -2:], out[0, 2, 1, -2:, -2:])
assert out.shape == t.Size((3, 4, 2, 4, 4))
return None
return None
def test_apply_jones_matrix_no_mult_modes_one_pattern_probe_mult_patterns_jones_1():
'''
probe:
2xMxL = 2x3x4
jones_matrix:
Nx2x2 = 3x2x2
jones matrices:
1: quat plate, 2: 90, 3: 0
'''
jones_matr = t.stack((jones_plate, jones90, jones0))
probe = t.ones(2, 3, 4, dtype=t.cfloat)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
print('expected shape: (3, 2, 3, 4)')
print('actual shape:', out.shape)
print('simulated:', out)
# we'll be checking only one mode (3th)
assert np.allclose(np.real(out[0, 0, :, :]), np.imag(out[0, 1, :, :]))
assert t.allclose(out[1, 0, :, :], t.zeros(3, 4, dtype=t.cfloat))
assert t.allclose(out[1, 1, :, :], t.ones(3, 4, dtype=t.cfloat))
assert t.allclose(out[2, 0, :, :], t.ones(3, 4, dtype=t.cfloat))
assert t.allclose(out[2, 1, :, :], t.zeros(3, 4, dtype=t.cfloat))
assert out.shape == t.Size((3, 2, 3, 4))
return None
def test_apply_jones_matrix_no_mult_modes_one_pattern_probe_mult_patterns_jones_2():
'''
probe:
2xMxL = 2x4x4
jones_matrix: jones matrices differ from pixel to pixel
Nx2x2xMxL = 3x2x2x4x4
jones matrix for the 1st pattern:
1: quat plate, 2: 90, 3: 0, 4: 45
jones matrix for the 2nd pattern:
1: 0, 2: 45, 3: 90, 4: quat plate
jones matrix for the 3rd pattern:
1: 90, 2: 0, 3: 45, 4: quat plate
'''
jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45),
build_from_quarters(jones0, jones45, jones90, jones_plate),
build_from_quarters(jones90, jones0, jones45, jones_plate)]
jones_matr = t.stack(([i for i in jones_m]), dim=0)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
print('expected shape: (3, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated:', out)
assert np.allclose(np.real(out[0, 0, :-2, :-2]), np.imag(out[0, 1, :-2, :-2]))
assert t.allclose(out[0, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 1, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 0, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:])
assert t.allclose(out[1, 0, :-2, :-2], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 0, :-2, -2:], out[1, 1, :-2, -2:])
assert t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
assert np.allclose(np.real(out[1, 0, -2:, -2:]), np.real(out[1, 0, -2:, -2:]))
assert t.allclose(out[2, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 1, :-2, :-2], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 0, -2:, :-2], out[2, 1, -2:, :-2])
assert np.allclose(np.real(out[2, 0, -2:, -2:]), np.real(out[2, 0, -2:, -2:]))
assert out.shape == t.Size((3, 2, 4, 4))
def test_apply_jones_matrix_mult_modes_one_pattern_probe_mult_patterns_jones_1():
'''
probe: multiple modes, multiple diffr patterns
Px2xMxL = 7x2x3x4
jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern)
Nx2x2 = 3x2x2
3 different matrices for each pattern:
1: 0
2: 45
3: 90
'''
probe = t.ones(2, 3, 4, dtype=t.cfloat)
probe = t.stack(([probe * 10 ** i for i in range(7)]), dim=0)
jones_matr = t.stack(([polarizer(angle) for angle in [0, 45, 90]]), dim=0)
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
print('expected shape: (3, 7, 2, 3, 4)')
print('actual:', out.shape)
print('simulated (patterns in one mode):', out[:, 6, :, :, :])
# we'll be checking only one mode
assert t.allclose(out[0, 6, 0, :, :], (10**6) * t.ones(3, 4, dtype=t.cfloat))
assert t.allclose(out[0, 6, 1, :, :], t.zeros(3, 4, dtype=t.cfloat))
assert t.allclose(out[1, 6, 0, :, :], out[1, 6, 1, :, :])
assert t.allclose(out[2, 6, 0, :, :], (10**6) * t.zeros(3, 4, dtype=t.cfloat))
assert t.allclose(out[2, 6, 1, :, :], (10**6) * t.ones(3, 4, dtype=t.cfloat))
assert out.shape == t.Size((3, 7, 2, 3, 4))
def test_apply_jones_matrix_mult_modes_one_pattern_probe_mult_patterns_jones_2():
'''
probe: multiple modes, multiple diffr pattern
Px2xMxL = 4x2x4x4
jones_matrix: jones matrices differ from pixel to pixel
Nx2x2xMxL = 3x2x2x4x4
(differs across the patterns in each mode)
jones matrix for the 1st pattern:
1: quat plate, 2: 90, 3: 0, 4: 45
jones matrix for the 2nd pattern:
1: 0, 2: 45, 3: 90, 4: quat plate
jones matrix for the 3rd pattern:
1: 90, 2: 0, 3: 45, 4: quat plate
'''
jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45),
build_from_quarters(jones0, jones45, jones90, jones_plate),
build_from_quarters(jones90, jones0, jones45, jones_plate)]
jones_matr = t.stack(([i for i in jones_m]), dim=0)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
probe = t.stack(([probe * (10 ** i) for i in range(4)]), dim=0)
print('probe:', probe.shape)
print('jones:', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
print('expected shape: (3, 4, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated patterns in one mode:', out[:, 3, :, :, :])
o = 10 ** 3
# we'll be checking only one mode (4th)
assert np.allclose(np.real(out[0, 3, 0, :-2, :-2]), np.imag(out[0, 3, 1, :-2, :-2]))
assert t.allclose(out[0, 3, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 3, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 3, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 3, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 3, 0, -2:, -2:], out[0, 3, 1, -2:, -2:])
assert t.allclose(out[1, 3, 0, :-2, :-2], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 3, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 3, 0, :-2, -2:], out[1, 3, 1, :-2, -2:])
assert t.allclose(out[1, 3, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 3, 1, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
assert np.allclose(np.real(out[1, 3, 0, -2:, -2:]), np.real(out[1, 3, 0, -2:, -2:]))
assert t.allclose(out[2, 3, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 3, 1, :-2, :-2], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 3, 0, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 3, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 3, 0, -2:, :-2], out[2, 3, 1, -2:, :-2])
assert np.allclose(np.real(out[2, 3, 0, -2:, -2:]), np.real(out[2, 3, 0, -2:, -2:]))
assert out.shape == t.Size((3, 4, 2, 4, 4))