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

Remove the original, worse unified mode model, and update the plotting functions to handle stacks of images, very helpful for multislice and also multi-mode models

Browse Source
This commit is contained in:
Abe Levitan
2021-02-04 18:55:55 -05:00
parent fdca857046
commit d664bf351b
9 changed files with 330 additions and 790 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/models/__init__.py
-1
View File
@@ -35,4 +35,3 @@ from CDTools.models.s_matrix_ptycho import SMatrixPtycho
from CDTools.models.multislice_2d_ptycho import Multislice2DPtycho
from CDTools.models.rpi import RPI
from CDTools.models.unified_mode_ptycho import UnifiedModePtycho
from CDTools.models.unified_mode_ptycho2 import UnifiedModePtycho2
CDTools/models/base.py
+1 -1
View File
@@ -311,7 +311,7 @@ class CDIModel(t.nn.Module):
# Define the optimizer
optimizer = t.optim.LBFGS(self.parameters(),
lr = lr, history_size=history_size)
return self.AD_optimize(iterations, data_loader, optimizer,
regularization_factor=regularization_factor,
thread=thread,
CDTools/models/fancy_ptycho.py
+5 -10
View File
@@ -153,6 +153,7 @@ class FancyPtycho(CDIModel):
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
else:
@@ -336,16 +337,10 @@ class FancyPtycho(CDIModel):
# Needs to be updated to allow for plotting to an existing figure
plot_list = [
('Dominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)),
('Dominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)),
('Subdominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis),
lambda self: len(self.probe) >=2),
('Subdominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis),
lambda self: len(self.probe) >=2),
('Probe Amplitude (scroll to view modes)',
lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)),
('Probe Phase (scroll to view modes)',
lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis)),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
('Object Phase',
CDTools/models/multislice_2d_ptycho.py
+20 -14
View File
@@ -387,26 +387,32 @@ class Multislice2DPtycho(CDIModel):
# Needs to be updated to allow for plotting to an existing figure
plot_list = [
('Dominant Probe Fourier Space Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[0] if self.fourier_probe else tools.propagators.inverse_far_field(self.probe[0]), fig=fig)),
('Dominant Probe Fourier Space Phase',
lambda self, fig: p.plot_phase(self.probe[0] if self.fourier_probe else tools.propagators.inverse_far_field(self.probe[0]), fig=fig)),
('Dominant Probe Real Space Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[0] if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe[0]), fig=fig, basis=self.probe_basis, units=self.units)),
('Dominant Probe Real Space Phase',
lambda self, fig: p.plot_phase(self.probe[0] if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe[0]), fig=fig, basis=self.probe_basis, units=self.units)),
('Subdominant Probe Real Space Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[1] if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe[1]), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: len(self.probe) >=2),
('Subdominant Probe Real Space Phase',
lambda self, fig: p.plot_phase(self.probe[1] if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe[1]), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: len(self.probe) >=2),
('Probe Fourier Space Amplitude',
lambda self, fig: p.plot_amplitude(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)),
('Probe Fourier Space Phase',
lambda self, fig: p.plot_phase(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)),
('Probe Real Space Amplitude',
lambda self, fig: p.plot_amplitude(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
('Probe Real Space Phase',
lambda self, fig: p.plot_phase(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
('Slice by Slice Real Part of T',
lambda self, fig: p.plot_real(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: self.exponentiate_obj),
('Slice by Slice Imaginary Part of T',
lambda self, fig: p.plot_imag(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: self.exponentiate_obj),
('Integrated Real Part of T',
lambda self, fig: p.plot_real(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: self.exponentiate_obj),
('Integrated Imaginary Part of T',
lambda self, fig: p.plot_imag(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: self.exponentiate_obj),
('Slice by Slice Amplitude of Object Function',
lambda self, fig: p.plot_amplitude(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: not self.exponentiate_obj),
('Slice by Slice Phase of Object Function',
lambda self, fig: p.plot_phase(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: not self.exponentiate_obj),
('Amplitude of Stacked Object Function',
lambda self, fig: p.plot_amplitude(reduce(cmath.cmult, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: not self.exponentiate_obj),
CDTools/models/unified_mode_ptycho.py
+122 -83
View File
@@ -18,7 +18,7 @@ class UnifiedModePtycho(CDIModel):
def __init__(self, wavelength, detector_geometry,
probe_basis,
probe_guess, obj_guess, rhos_guess,
probe_guess, obj_guess, Ws_guess,
detector_slice=None,
surface_normal=np.array([0.,0.,1.]),
min_translation = t.Tensor([0,0]),
@@ -71,7 +71,10 @@ class UnifiedModePtycho(CDIModel):
self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
self.rhos = t.nn.Parameter(t.Tensor(rhos_guess).to(t.float32))
if type(Ws_guess) == type(t.zeros(1)):
self.Ws = t.nn.Parameter(Ws_guess.to(t.float32))
else:
self.Ws = t.nn.Parameter(cmath.complex_to_torch(Ws_guess).to(t.float32))
if translation_offsets is None:
self.translation_offsets = None
@@ -95,7 +98,7 @@ class UnifiedModePtycho(CDIModel):
@classmethod
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, opt_for_fft=False, mixing_mode='unified'):
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, opt_for_fft=False, dm_rank=0):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -175,13 +178,22 @@ class UnifiedModePtycho(CDIModel):
translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
#
if mixing_mode.lower().strip() == 'unified':
rhos = t.zeros(len(dataset),n_modes,n_modes)
rhos[:,0,0] = 1
for i in range(1,n_modes):
rhos[:,i,i] = 1/n_modes
# dm_rank defines the rank of the shot-by-shot density matrices
if dm_rank > n_modes:
raise KeyError('Density matrix rank cannot be greater than the number of modes')
elif dm_rank != 0:
if dm_rank == -1:
dm_rank = n_modes
Ws = t.zeros(len(dataset),dm_rank,n_modes,2)
Ws[:,0,0,0] = 1
for i in range(1,dm_rank):
Ws[:,i,i,0] = 1/np.sqrt(n_modes)
else:
# dm_rank=0 is a special case defining a purely stable, incoherent
# mode mixing model. This is passed on by defining a set of weights
# which only has one index
Ws = t.ones(len(dataset))
if hasattr(dataset, 'mask') and dataset.mask is not None:
mask = dataset.mask.to(t.bool)
else:
@@ -207,7 +219,7 @@ class UnifiedModePtycho(CDIModel):
else:
obj_support = None
return cls(wavelength, det_geo, probe_basis, probe, obj, rhos,
return cls(wavelength, det_geo, probe_basis, probe, obj, Ws,
detector_slice=det_slice,
surface_normal=surface_normal,
min_translation=min_translation,
@@ -229,22 +241,61 @@ class UnifiedModePtycho(CDIModel):
if self.translation_offsets is not None:
pix_trans += self.translation_scale * self.translation_offsets[index]
Ws = self.Ws[index]
# This probably needs to be fixed, I doubt it will really still work
# to deal with single-pattern sims.
#if type(index) == type(0):
# print('hi')
# index = [index]
# Ws = [Ws]
# pix_trans = [pix_trans]
# single_frame = True
#else:
# single_frame = False
probes = []
all_exit_waves = []
for i in range(self.probe.shape[0]):
# from storing the probe in Fourier space
#pr = tools.propagators.inverse_far_field(self.probe[i]) * self.probe_support
pr = self.probe[i] * self.probe_support
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(pr,
self.obj_support * self.obj,
pix_trans,
shift_probe=True)
exit_waves = exit_waves * self.probe_support[...,:,:]
all_exit_waves.append(exit_waves)
# This is the case if a purely stable, incoherent model is defined.
#if len(Ws[0].shape) == 0:
# What we do here is generate an identity matrix, and multiply
# that identity matrix by the per-frame weight.
#Ws = [W * t.stack([t.eye(self.probe.shape[0]),
# t.zeros([self.probe.shape[0]]*2)],dim=-1).to(
# dtype=W.dtype, device=W.device)
# for W in Ws]
#print(Ws)
# This restricts the basis probes with the probe support
basis_prs = self.probe * self.probe_support[...,:,:]
return t.stack(all_exit_waves)
if len(Ws[0].shape) == 0:
# If a purely stable coherent illumination is defined
prs = cmath.cmult(Ws[...,None,None,None,:],basis_prs)
else:
# If a frame-by-frame weight matrix is defined
# This takes the dot product of all the weight matrices with
# the probes. The output has dimensions of translation, then
# coherent mode index, then x,y, and then complex index
prs = t.sum(cmath.cmult(Ws[...,None,None,:], basis_prs),
axis=-4)
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs, self.obj_support * self.obj,pix_trans,
shift_probe=True, multiple_modes=True)
exit_waves = exit_waves * self.probe_support[...,:,:]
if hasattr(self,'weights') and self.weights is not None:
if exit_waves.dim() == 5:
exit_waves = self.weights[index][:,None,None,None,None] \
* exit_waves
else:
exit_waves = self.weights[index] * exit_waves
return exit_waves
def forward_propagator(self, wavefields):
@@ -255,28 +306,14 @@ class UnifiedModePtycho(CDIModel):
return tools.propagators.inverse_far_field(wavefields)
def measurement(self, wavefields, indices):
#return tools.measurements.density_matrix(wavefields,self.rhos[indices],
# detector_slice=self.detector_slice,
# saturation=self.saturation,
# oversampling=self.oversampling)
def measurement(self, wavefields):
return tools.measurements.quadratic_background(wavefields,
self.background, self.rhos[indices],
self.background,
detector_slice=self.detector_slice,
measurement=tools.measurements.density_matrix,
measurement=tools.measurements.incoherent_sum,
saturation=self.saturation,
oversampling=self.oversampling)
def forward(self, *args):
"""The complete forward model
We need to override this to enable the wavefield mixing at the
level of the measurement function
"""
indices = args[0]
return self.measurement(self.forward_propagator(self.interaction(*args)),indices)
def loss(self, sim_data, real_data, mask=None):
return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
#return tools.losses.poisson_nll(real_data, sim_data, mask=mask)
@@ -344,22 +381,15 @@ class UnifiedModePtycho(CDIModel):
mask=mask)
def get_rhos(self):
rhos_out = np.zeros([self.rhos.shape[0],
self.rhos.shape[1],self.rhos.shape[2]],
# If this is not a purely stable model
if len(self.Ws.shape) >= 2:
Ws = cmath.torch_to_complex(self.Ws.detach().cpu())
rhos_out = np.matmul(np.swapaxes(Ws,1,2), Ws.conj())
return rhos_out
else:
return np.array([np.eye(self.probe.shape[0])]*self.Ws.shape[0],
dtype=np.complex64)
for (i,j) in ((i,j) for i in range(rhos_out.shape[1])
for j in range(rhos_out.shape[2])):
if i == j:
rhos_out[:,i,j] += self.rhos.data[:,i,j].cpu().detach().numpy()
if i < j: # upper triangle, real part
rhos_out[:,i,j] += self.rhos.data[:,i,j].cpu().detach().numpy()
rhos_out[:,j,i] += self.rhos.data[:,i,j].cpu().detach().numpy()
if i > j: # upper triangle, real part
rhos_out[:,j,i] += 1j * self.rhos.data[:,i,j].cpu().detach().numpy()
rhos_out[:,i,j] -= 1j * self.rhos.data[:,i,j].cpu().detach().numpy()
return rhos_out
def tidy_probes(self, normalization=1):
"""Tidies up the probes
@@ -368,6 +398,18 @@ class UnifiedModePtycho(CDIModel):
density matrices to operate in that updated basis
"""
# Must also implement a version that works appropriately with
# a purely incoherent model
#
# Note to future: We could probably do this more cleanly with an
# SVD directly on the Ws matrix, instead of an eigendecomposition
# of the rho matrix. This could avoid potential stability issues
# due to the existence of zero eigenvalues in the full rho matrix
# when dm_rank < n_modes
#
rhos = self.get_rhos()
overall_rho = np.mean(rhos,axis=0)
probe = cmath.torch_to_complex(self.probe.detach().cpu())
@@ -377,24 +419,26 @@ class UnifiedModePtycho(CDIModel):
normalize=True)
Aconj = A.conj()
Atrans = np.transpose(A)
new_rhos = np.swapaxes(np.dot(Atrans,np.dot(rhos,Aconj)),0,1)
new_rhos = np.matmul(Atrans,np.matmul(rhos,Aconj))
new_rhos /= normalization
ortho_probes *= np.sqrt(normalization)
#print(np.dot(Aconjinv,rhos).shape)
#print(np.dot(rhos,Aconj).shape)
new_rhos = cmath.complex_to_torch(new_rhos).to(
dtype=self.rhos.dtype,device=self.rhos.device)
# This repacks the data into the format used internally
for (i,j) in ((i,j) for i in range(new_rhos.shape[1])
for j in range(new_rhos.shape[2])):
if i == j:
self.rhos.data[:,i,j] = new_rhos[:,i,j,0]
if i < j: # upper triangle, real part
self.rhos.data[:,i,j] = new_rhos[:,i,j,0]
self.rhos.data[:,j,i] = new_rhos[:,i,j,1]
dm_rank = self.Ws.shape[1]
new_Ws = []
for rho in new_rhos:
# These are returned from smalles to largest - we want to keep
# the largest ones
w,v = np.linalg.eigh(rho)
w = w[::-1][:dm_rank]
v = v[:,::-1][:,:dm_rank]
new_Ws.append(np.dot(np.diag(np.sqrt(w)),v.transpose()))
new_Ws = np.array(new_Ws)
self.Ws.data = cmath.complex_to_torch(new_Ws).to(
dtype=self.Ws.dtype,device=self.Ws.device)
self.probe.data = cmath.complex_to_torch(ortho_probes).to(
device=self.probe.device,dtype=self.probe.dtype)
@@ -408,21 +452,16 @@ class UnifiedModePtycho(CDIModel):
# Needs to be updated to allow for plotting to an existing figure
plot_list = [
('Dominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)),
('Dominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)),
('Subdominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis),
lambda self: len(self.probe) >=2),
('Subdominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis),
lambda self: len(self.probe) >=2),
('Basis Probe Amplitudes',
lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)),
('Basis Probe Phases',
lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis)),
('Average Density Matrix Amplitudes',
lambda self, fig: p.plot_amplitude(np.mean(np.abs(self.get_rhos()),axis=0), fig=fig),
lambda self: len(self.probe) >=2),
('Von Neumann Entropy (only accurate after tidy_probes)',
lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_vn_entropy(self.get_rhos()), fig=fig)),
lambda self: len(self.Ws.shape) >=2),
('% Power in Top Mode (only accurate after tidy_probes)',
lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_top_mode_fraction(self.get_rhos()), fig=fig),
lambda self: len(self.Ws.shape) >=2),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
('Object Phase',
@@ -441,9 +480,9 @@ class UnifiedModePtycho(CDIModel):
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = cmath.torch_to_complex(self.obj.detach().cpu())
background = self.background.detach().cpu().numpy()**2
weights = self.weights.detach().cpu().numpy()
Ws = cmath.torch_to_complex(self.Ws.detach().cpu())
return {'basis':basis, 'translation':translations,
'probe':probe,'obj':obj,
'background':background,
'weights':weights}
'Ws':Ws}
CDTools/models/unified_mode_ptycho2.py
-490
View File
@@ -1,490 +0,0 @@
from __future__ import division, print_function, absolute_import
import torch as t
from CDTools.models import CDIModel
from CDTools.datasets import Ptycho2DDataset
from CDTools import tools
from CDTools.tools import cmath
from CDTools.tools import analysis
from CDTools.tools import plotting as p
from matplotlib import pyplot as plt
from datetime import datetime
import numpy as np
from copy import copy
__all__ = ['UnifiedModePtycho2']
class UnifiedModePtycho2(CDIModel):
def __init__(self, wavelength, detector_geometry,
probe_basis,
probe_guess, obj_guess, Ws_guess,
detector_slice=None,
surface_normal=np.array([0.,0.,1.]),
min_translation = t.Tensor([0,0]),
background = None, translation_offsets=None, mask=None,
translation_scale = 1, saturation=None,
probe_support = None, obj_support=None, oversampling=1):
super(UnifiedModePtycho2,self).__init__()
self.wavelength = t.Tensor([wavelength])
self.detector_geometry = copy(detector_geometry)
det_geo = self.detector_geometry
if hasattr(det_geo, 'distance'):
det_geo['distance'] = t.Tensor(det_geo['distance'])
if hasattr(det_geo, 'basis'):
det_geo['basis'] = t.Tensor(det_geo['basis'])
if hasattr(det_geo, 'corner'):
det_geo['corner'] = t.Tensor(det_geo['corner'])
self.min_translation = t.Tensor(min_translation)
self.probe_basis = t.Tensor(probe_basis)
self.detector_slice = detector_slice
self.surface_normal = t.Tensor(surface_normal)
self.saturation = saturation
if mask is None:
self.mask = mask
else:
self.mask = t.BoolTensor(mask)
# We rescale the probe here so it learns at the same rate as the
# object
if probe_guess.dim() > 3:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32)))
else:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32)))
self.probe = t.nn.Parameter(probe_guess.to(t.float32)
/ self.probe_norm)
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
if background is None:
if detector_slice is not None:
background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1])
else:
background = 1e-6 * t.ones(self.probe[0].shape[:-1])
self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
if type(Ws_guess) == type(t.zeros(1)):
self.Ws = t.nn.Parameter(Ws_guess.to(t.float32))
else:
self.Ws = t.nn.Parameter(cmath.complex_to_torch(Ws_guess).to(t.float32))
if translation_offsets is None:
self.translation_offsets = None
else:
self.translation_offsets = t.nn.Parameter(t.Tensor(translation_offsets).to(t.float32)/ translation_scale)
self.translation_scale = translation_scale
if probe_support is not None:
self.probe_support = probe_support
else:
self.probe_support = t.ones_like(self.probe[0])
if obj_support is not None:
self.obj_support = obj_support
self.obj.data = self.obj * obj_support
else:
self.obj_support = t.ones_like(self.obj)
self.oversampling = oversampling
@classmethod
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, opt_for_fft=False, dm_rank=0):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
det_shape = dataset[0][1].shape
distance = dataset.detector_geometry['distance']
# always do this on the cpu
get_as_args = dataset.get_as_args
dataset.get_as(device='cpu')
(indices, translations), patterns = dataset[:]
dataset.get_as(*get_as_args[0],**get_as_args[1])
# Set to none to avoid issues with things outside the detector
if auto_center:
center = tools.image_processing.centroid(t.sum(patterns,dim=0))
else:
center = None
# Then, generate the probe geometry from the dataset
ewg = tools.initializers.exit_wave_geometry
probe_basis, probe_shape, det_slice = ewg(det_basis,
det_shape,
wavelength,
distance,
center=center,
padding=padding,
opt_for_fft=opt_for_fft,
oversampling=oversampling)
if hasattr(dataset, 'sample_info') and \
dataset.sample_info is not None and \
'orientation' in dataset.sample_info:
surface_normal = dataset.sample_info['orientation'][2]
else:
surface_normal = np.array([0.,0.,1.])
# If this information is supplied when the function is called,
# then we override the information in the .cxi file
if scattering_mode in {'t', 'transmission'}:
surface_normal = np.array([0.,0.,1.])
elif scattering_mode in {'r', 'reflection'}:
outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1])
outgoing_dir /= np.linalg.norm(outgoing_dir)
surface_normal = outgoing_dir + np.array([0.,0.,1.])
surface_normal /= -np.linalg.norm(surface_normal)
# Next generate the object geometry from the probe geometry and
# the translations
pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations, surface_normal=surface_normal)
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
else:
background = None
# Finally, initialize the probe and object using this information
if probe_size is None:
probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
else:
probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
# Now we initialize all the subdominant probe modes
probe_max = t.max(cmath.cabs(probe))
probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)]
probe = t.stack([probe,] + probe_stack)
#probe = t.stack([tools.propagators.far_field(probe),] + probe_stack)
obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
det_geo = dataset.detector_geometry
translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
# dm_rank defines the rank of the shot-by-shot density matrices
if dm_rank > n_modes:
raise KeyError('Density matrix rank cannot be greater than the number of modes')
elif dm_rank != 0:
if dm_rank == -1:
dm_rank = n_modes
Ws = t.zeros(len(dataset),dm_rank,n_modes,2)
Ws[:,0,0,0] = 1
for i in range(1,dm_rank):
Ws[:,i,i,0] = 1/np.sqrt(n_modes)
else:
# dm_rank=0 is a special case defining a purely stable, incoherent
# mode mixing model. This is passed on by defining a set of weights
# which only has one index
Ws = t.ones(len(dataset))
if hasattr(dataset, 'mask') and dataset.mask is not None:
mask = dataset.mask.to(t.bool)
else:
mask = None
if probe_support_radius is not None:
probe_support = t.zeros_like(probe[0].to(dtype=t.float32))
p_cent = np.array(probe.shape[1:3]).astype(int) // 2
psr = int(probe_support_radius)
probe_support[p_cent[0]-psr:p_cent[0]+psr,
p_cent[1]-psr:p_cent[1]+psr] = 1
probe = probe * probe_support[None,:,:]
else:
probe_support = None;
if restrict_obj != -1:
ro = restrict_obj
os = np.array(obj_size)
ps = np.array(probe_shape)
obj_support = t.zeros_like(obj.to(dtype=t.float32))
obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2,
ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1
else:
obj_support = None
return cls(wavelength, det_geo, probe_basis, probe, obj, Ws,
detector_slice=det_slice,
surface_normal=surface_normal,
min_translation=min_translation,
translation_offsets = translation_offsets,
mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
probe_support=probe_support,
obj_support=obj_support,
oversampling=oversampling)
def interaction(self, index, translations):
pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
translations,
surface_normal=self.surface_normal)
pix_trans -= self.min_translation
if self.translation_offsets is not None:
pix_trans += self.translation_scale * self.translation_offsets[index]
Ws = self.Ws[index]
if type(index) == type(0):
index = [index]
Ws = [Ws]
pix_trans = [pix_trans]
single_frame = True
else:
single_frame = False
probes = []
all_exit_waves = []
# This is the case if a purely stable, incoherent model is defined.
if len(Ws[0].shape) == 0:
# What we do here is generate an identity matrix, and multiply
# that identity matrix by the per-frame weight.
Ws = [W * t.stack([t.eye(self.probe.shape[0]),
t.zeros([self.probe.shape[0]]*2)],dim=-1).to(
dtype=W.dtype, device=W.device)
for W in Ws]
# Outer iteration is the mode index iteration
for i in range(Ws[0].shape[0]):
# Now we need to separately treat each mode
exit_waves = []
for W, pix_tran in zip(Ws, pix_trans):
# from storing the probe in Fourier space
pr = [cmath.cmult(W[i,j,:], self.probe[j] * self.probe_support)
for j in range(self.probe.shape[0])]
pr = t.sum(t.stack(pr), axis=0)
exit_waves.append(self.probe_norm *
tools.interactions.ptycho_2D_sinc(
pr, self.obj_support * self.obj,
pix_tran, shift_probe=True))
exit_waves = t.stack(exit_waves)
if single_frame:
exit_waves = exit_waves[0]
# Multiply again by probe support to suppress the fringes from the
# sinc-interpolated shift
exit_waves = exit_waves * self.probe_support[...,:,:]
all_exit_waves.append(exit_waves)
return t.stack(all_exit_waves)
def forward_propagator(self, wavefields):
return tools.propagators.far_field(wavefields)
def backward_propagator(self, wavefields):
return tools.propagators.inverse_far_field(wavefields)
def measurement(self, wavefields):
return tools.measurements.quadratic_background(wavefields,
self.background,
detector_slice=self.detector_slice,
measurement=tools.measurements.incoherent_sum,
saturation=self.saturation,
oversampling=self.oversampling)
def loss(self, sim_data, real_data, mask=None):
return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
#return tools.losses.poisson_nll(real_data, sim_data, mask=mask)
def to(self, *args, **kwargs):
super(UnifiedModePtycho2, self).to(*args, **kwargs)
self.wavelength = self.wavelength.to(*args,**kwargs)
# move the detector geometry too
det_geo = self.detector_geometry
if hasattr(det_geo, 'distance'):
det_geo['distance'] = det_geo['distance'].to(*args,**kwargs)
if hasattr(det_geo, 'basis'):
det_geo['basis'] = det_geo['basis'].to(*args,**kwargs)
if hasattr(det_geo, 'corner'):
det_geo['corner'] = det_geo['corner'].to(*args,**kwargs)
if self.mask is not None:
self.mask = self.mask.to(*args, **kwargs)
self.min_translation = self.min_translation.to(*args,**kwargs)
self.probe_basis = self.probe_basis.to(*args,**kwargs)
self.probe_norm = self.probe_norm.to(*args,**kwargs)
self.probe_support = self.probe_support.to(*args,**kwargs)
self.obj_support = self.obj_support.to(*args,**kwargs)
self.surface_normal = self.surface_normal.to(*args, **kwargs)
def sim_to_dataset(self, args_list):
# In the future, potentially add more control
# over what metadata is saved (names, etc.)
# First, I need to gather all the relevant data
# that needs to be added to the dataset
entry_info = {'program_name': 'CDTools',
'instrument_n': 'Simulated Data',
'start_time': datetime.now()}
surface_normal = self.surface_normal.detach().cpu().numpy()
xsurfacevec = np.cross(np.array([0.,1.,0.]), surface_normal)
xsurfacevec /= np.linalg.norm(xsurfacevec)
ysurfacevec = np.cross(surface_normal, xsurfacevec)
ysurfacevec /= np.linalg.norm(ysurfacevec)
orientation = np.array([xsurfacevec, ysurfacevec, surface_normal])
sample_info = {'description': 'A simulated sample',
'orientation': orientation}
detector_geometry = self.detector_geometry
mask = self.mask
wavelength = self.wavelength
indices, translations = args_list
# Then we simulate the results
data = self.forward(indices, translations)
# And finally, we make the dataset
return Ptycho2DDataset(translations, data,
entry_info = entry_info,
sample_info = sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask)
def get_rhos(self):
# If this is not a purely stable model
if len(self.Ws.shape) >= 2:
Ws = cmath.torch_to_complex(self.Ws.detach().cpu())
rhos_out = np.matmul(np.swapaxes(Ws,1,2), Ws.conj())
return rhos_out
else:
return np.array([np.eye(self.probe.shape[0])]*self.Ws.shape[0],
dtype=np.complex64)
def tidy_probes(self, normalization=1):
"""Tidies up the probes
What we want to do here is use all the information on all the probes
to calculate a natural basis for the experiment, and update all the
density matrices to operate in that updated basis
"""
# Must also implement a version that works appropriately with
# a purely incoherent model
#
# Note to future: We could probably do this more cleanly with an
# SVD directly on the Ws matrix, instead of an eigendecomposition
# of the rho matrix. This could avoid potential stability issues
# due to the existence of zero eigenvalues in the full rho matrix
# when dm_rank < n_modes
#
rhos = self.get_rhos()
overall_rho = np.mean(rhos,axis=0)
probe = cmath.torch_to_complex(self.probe.detach().cpu())
ortho_probes, A = analysis.orthogonalize_probes(probe,
density_matrix=overall_rho,
keep_transform=True,
normalize=True)
Aconj = A.conj()
Atrans = np.transpose(A)
new_rhos = np.matmul(Atrans,np.matmul(rhos,Aconj))
new_rhos /= normalization
ortho_probes *= np.sqrt(normalization)
dm_rank = self.Ws.shape[1]
new_Ws = []
for rho in new_rhos:
# These are returned from smalles to largest - we want to keep
# the largest ones
w,v = np.linalg.eigh(rho)
w = w[::-1][:dm_rank]
v = v[:,::-1][:,:dm_rank]
new_Ws.append(np.dot(np.diag(np.sqrt(w)),v.transpose()))
new_Ws = np.array(new_Ws)
self.Ws.data = cmath.complex_to_torch(new_Ws).to(
dtype=self.Ws.dtype,device=self.Ws.device)
self.probe.data = cmath.complex_to_torch(ortho_probes).to(
device=self.probe.device,dtype=self.probe.dtype)
def corrected_translations(self,dataset):
translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device)
t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
return translations + t_offset
# Needs to be updated to allow for plotting to an existing figure
plot_list = [
('Dominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)),
('Dominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)),
('Subdominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis),
lambda self: len(self.probe) >=2),
('Subdominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis),
lambda self: len(self.probe) >=2),
('Average Density Matrix Amplitudes',
lambda self, fig: p.plot_amplitude(np.mean(np.abs(self.get_rhos()),axis=0), fig=fig),
lambda self: len(self.Ws.shape) >=2),
('% Power in Top Mode (only accurate after tidy_probes)',
lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_top_mode_fraction(self.get_rhos()), fig=fig),
lambda self: len(self.Ws.shape) >=2),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
('Object Phase',
lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
('Background',
lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
]
def save_results(self, dataset):
basis = self.probe_basis.detach().cpu().numpy()
translations = self.corrected_translations(dataset).detach().cpu().numpy()
probe = cmath.torch_to_complex(self.probe.detach().cpu())
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = cmath.torch_to_complex(self.obj.detach().cpu())
background = self.background.detach().cpu().numpy()**2
Ws = cmath.torch_to_complex(self.Ws.detach().cpu())
return {'basis':basis, 'translation':translations,
'probe':probe,'obj':obj,
'background':background,
'Ws':Ws}
CDTools/tools/plotting/plotting.py
+179 -187
View File
@@ -82,7 +82,146 @@ def get_units_factor(units):
return factor
def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs):
def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label=None, **kwargs):
"""Plots an image with a colorbar and on an appropriate spatial grid
If a figure is given explicitly, it will clear that existing figure and
plot over it. Otherwise, it will generate a new figure.
If a basis is explicitly passed, the image will be plotted in real-space
coordinates
Finally, if a function is passed to the plot_func argument, this function
will be called on each slice of data before it is plotted. This is used
internally to enable the plot_real, plot_image, plot_phase, etc. functions.
Parameters
----------
im : array
An complex array with dimensions NxM
plot_func : callable
A function which maps numpy arrays to the image to be plotted
fig : matplotlib.figure.Figure
Default is a new figure, a matplotlib figure to use to plot
basis : np.array
Optional, the 3x2 probe basis
units : str
The length units to mark on the plot, default is um
cmap : str
Default is 'viridis', the colormap to plot with
cmap_label : str
What to label the colorbar when plotting
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
Returns
-------
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
# convert to numpy
if isinstance(im, t.Tensor):
# If final dimension is 2, assume it is a complex array. If not,
# assume it represents a real array
if im.shape[-1] == 2:
im = cmath.torch_to_complex(im.detach().cpu())
else:
im = im.detach().cpu().numpy()
if fig is None:
fig = plt.figure()
ax = fig.add_subplot(111, **kwargs)
# This nukes everything and updates either the appropriate image from the
# stack of images, or the only image if only a single image has been
# given
def make_plot(idx):
plt.figure(fig.number)
title = plt.gca().get_title()
fig.clear()
# If im only has two dimensions, this reshape will add a leading
# dimension, and update will be called on index 0. If it has 3 or more
# dimensions, then all the leading dimensions will be compressed into
# one long dimension which can be scrolled through.
s = im.shape
reshaped_im = im.reshape(-1,s[-2],s[-1])
num_images = reshaped_im.shape[0]
fig.plot_idx = idx % num_images
to_plot = plot_func(reshaped_im[fig.plot_idx])
#Plot in a basis if it exists, otherwise dont
if basis is not None:
if isinstance(basis,t.Tensor):
np_basis = basis.detach().cpu().numpy()
else:
np_basis = basis
# This fails if the basis is not rectangular
basis_norm = np.linalg.norm(np_basis, axis = 0)
basis_norm = basis_norm * get_units_factor(units)
extent = [0, to_plot.shape[-1]*basis_norm[1], 0,
to_plot.shape[-2]*basis_norm[0]]
else:
extent=None
plt.imshow(to_plot, cmap = cmap, extent = extent)
cbar = plt.colorbar()
if cmap_label is not None:
cbar.set_label(cmap_label)
if basis is not None:
plt.xlabel('X (' + units + ')')
plt.ylabel('Y (' + units + ')')
else:
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
plt.title(title)
if len(im.shape) >= 3:
plt.text(0.03, 0.03, str(fig.plot_idx), fontsize=14, transform=plt.gcf().transFigure)
return fig
if hasattr(fig, 'plot_idx'):
result = make_plot(fig.plot_idx)
else:
result = make_plot(0)
update = make_plot
def on_action(event):
if not hasattr(event, 'button'):
event.button = None
if not hasattr(event, 'key'):
event.key = None
if event.key == 'up' or event.button == 'up':
update(fig.plot_idx - 1)
elif event.key == 'down' or event.button == 'down':
update(fig.plot_idx + 1)
plt.draw()
if len(im.shape) >=3:
if not hasattr(fig,'my_callbacks'):
fig.my_callbacks = []
for cid in fig.my_callbacks:
fig.canvas.mpl_disconnect(cid)
fig.my_callbacks = []
fig.my_callbacks.append(fig.canvas.mpl_connect('key_press_event',on_action))
fig.my_callbacks.append(fig.canvas.mpl_connect('scroll_event',on_action))
return result
def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Real Part (a.u.)', **kwargs):
"""Plots the real part of a complex array with dimensions NxM
If a figure is given explicitly, it will clear that existing figure and
@@ -103,6 +242,8 @@ def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwa
The length units to mark on the plot, default is um
cmap : str
Default is 'viridis', the colormap to plot with
cmap_label : str
What to label the colorbar when plotting
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
@@ -111,45 +252,14 @@ def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwa
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
if fig is None:
fig = plt.figure()
ax = fig.add_subplot(111, **kwargs)
else:
plt.figure(fig.number)
plt.gcf().clear()
if isinstance(im, t.Tensor):
real = im[...,0].detach().cpu().numpy()
else:
real = np.real(im)
#Plot in a basis if it exists, otherwise dont
if basis is not None:
if isinstance(basis,t.Tensor):
basis = basis.detach().cpu().numpy()
# This fails if the
basis_norm = np.linalg.norm(basis, axis = 0)
basis_norm = basis_norm * get_units_factor(units)
extent = [0, real.shape[-1]*basis_norm[1], 0, real.shape[-2]*basis_norm[0]]
else:
extent=None
plt.imshow(real, cmap = cmap, extent = extent)
cbar = plt.colorbar()
cbar.set_label('Real Part (a.u.)')
if basis is not None:
plt.xlabel('X (' + units + ')')
plt.ylabel('Y (' + units + ')')
else:
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
return fig
plot_func = lambda x: np.real(x)
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, cmap=cmap, cmap_label=cmap_label,
**kwargs)
def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs):
def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Imaginary Part (a.u.)', **kwargs):
"""Plots the imaginary part of a complex array with dimensions NxM
If a figure is given explicitly, it will clear that existing figure and
@@ -170,6 +280,8 @@ def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwa
The length units to mark on the plot, default is um
cmap : str
Default is 'viridis', the colormap to plot with
cmap_label : str
What to label the colorbar when plotting
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
@@ -178,53 +290,21 @@ def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwa
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
if fig is None:
fig = plt.figure()
ax = fig.add_subplot(111, **kwargs)
else:
plt.figure(fig.number)
plt.gcf().clear()
if isinstance(im, t.Tensor):
imag = im[...,1].detach().cpu().numpy()
else:
imag = np.imag(im)
#Plot in a basis if it exists, otherwise dont
if basis is not None:
if isinstance(basis,t.Tensor):
basis = basis.detach().cpu().numpy()
# This fails if the
basis_norm = np.linalg.norm(basis, axis = 0)
basis_norm = basis_norm * get_units_factor(units)
extent = [0, imag.shape[-1]*basis_norm[1], 0, imag.shape[-2]*basis_norm[0]]
else:
extent=None
plt.imshow(imag, cmap = cmap, extent = extent)
cbar = plt.colorbar()
cbar.set_label('Imaginary Part (a.u.)')
if basis is not None:
plt.xlabel('X (' + units + ')')
plt.ylabel('Y (' + units + ')')
else:
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
return fig
plot_func = lambda x: np.imag(x)
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, cmap=cmap, cmap_label=cmap_label,
**kwargs)
def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs):
def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Amplitude (a.u.)', **kwargs):
"""Plots the amplitude of a complex array with dimensions NxM
If a figure is given explicitly, it will clear that existing figure and
plot over it. Otherwise, it will generate a new figure.
If a basis is explicitly passed, the image will be plotted in real-space
coordinates
coordinates.
Parameters
----------
im : array
@@ -237,6 +317,8 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis',
The length units to mark on the plot, default is um
cmap : str
Default is 'viridis', the colormap to plot with
cmap_label : str
What to label the colorbar when plotting
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
@@ -245,45 +327,13 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis',
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
if fig is None:
fig = plt.figure()
ax = fig.add_subplot(111, **kwargs)
else:
plt.figure(fig.number)
plt.gcf().clear()
if isinstance(im, t.Tensor):
absolute = cmath.cabs(im).detach().cpu().numpy()
else:
absolute = np.absolute(im)
#Plot in a basis if it exists, otherwise dont
if basis is not None:
if isinstance(basis,t.Tensor):
basis = basis.detach().cpu().numpy()
# This fails if the
basis_norm = np.linalg.norm(basis, axis = 0)
basis_norm = basis_norm * get_units_factor(units)
extent = [0, absolute.shape[-1]*basis_norm[1], 0, absolute.shape[-2]*basis_norm[0]]
else:
extent=None
plt.imshow(absolute, cmap = cmap, extent = extent)
cbar = plt.colorbar()
cbar.set_label('Amplitude (a.u.)')
if basis is not None:
plt.xlabel('X (' + units + ')')
plt.ylabel('Y (' + units + ')')
else:
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
return fig
plot_func = lambda x: np.absolute(x)
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, cmap=cmap, cmap_label=cmap_label,
**kwargs)
def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', **kwargs):
def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', cmap_label='Phase (rad)', **kwargs):
""" Plots the phase of a complex array with dimensions NxMx2
If a figure is given explicitly, it will clear that existing figure and
@@ -304,6 +354,8 @@ def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', **kwargs)
The length units to mark on the plot, default is um
cmap : str
Default is 'viridis', the colormap to plot with
cmap_label : str
What to label the colorbar when plotting
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
@@ -312,49 +364,18 @@ def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', **kwargs)
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
if fig is None:
fig = plt.figure()
ax = fig.add_subplot(111, **kwargs)
else:
plt.figure(fig.number)
plt.gcf().clear()
if isinstance(im, t.Tensor):
phase = cmath.cphase(im).detach().cpu().numpy()
else:
phase = np.angle(im)
if basis is not None:
if isinstance(basis,t.Tensor):
basis = basis.detach().cpu().numpy()
basis_norm = np.linalg.norm(basis, axis = 0)
basis_norm = basis_norm * get_units_factor(units)
extent = [0, phase.shape[-1]*basis_norm[1], 0, phase.shape[-2]*basis_norm[0]]
else:
extent=None
# If the user has matplotlib >=3.0, use the preferred colormap
if cmap == 'auto':
try:
plt.imshow(phase, cmap = 'twilight', extent=extent)
except:
plt.imshow(phase, cmap = 'hsv', extent=extent)
else:
plt.imshow(phase, cmap = cmap, extent=extent)
if 'twilight' in plt.colormaps():
cmap = 'twilight'
elif 'hsv' in plt.colormaps():
cmap = 'hsv'
else:
raise AttributeError('Neither twilight or hsv colormap exists in this screwed up matplotlib install')
cbar = plt.colorbar()
cbar.set_label('Phase (rad)')
if basis is not None:
plt.xlabel('X (' + units + ')')
plt.ylabel('Y (' + units + ')')
else:
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
return fig
plot_func = lambda x: np.angle(x)
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, cmap=cmap, cmap_label=cmap_label,
**kwargs)
def plot_amplitude_surfacenorm():
@@ -390,38 +411,9 @@ def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs):
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
if fig is None:
fig = plt.figure()
ax = fig.add_subplot(111, **kwargs)
else:
plt.figure(fig.number)
plt.gcf().clear()
if isinstance(im, t.Tensor):
im = cmath.torch_to_complex(im.detach().cpu())
if basis is not None:
if isinstance(basis,t.Tensor):
basis = basis.detach().cpu().numpy()
basis_norm = np.linalg.norm(basis, axis = 0)
basis_norm = basis_norm * get_units_factor(units)
extent = [0, im.shape[-1]*basis_norm[1], 0, im.shape[-2]*basis_norm[0]]
else:
extent=None
colorized = colorize(im)
plt.imshow(colorized, extent=extent)
if basis is not None:
plt.xlabel('X (' + units + ')')
plt.ylabel('Y (' + units + ')')
else:
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
return fig
plot_func = lambda x: colorize(x)
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, cmap=cmap, **kwargs)
def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwargs):
examples/gold_ball_ptycho.py
+2 -2
View File
@@ -17,8 +17,8 @@ model = CDTools.models.FancyPtycho.from_dataset(dataset, n_modes=2)
model.to(device='cuda')
dataset.get_as(device='cuda')
#for i, loss in enumerate(model.Adam_optimize(20, dataset, batch_size=50)):
for i, loss in enumerate(model.LBFGS_optimize(20, dataset, lr=1, history_size=5)):
for i, loss in enumerate(model.Adam_optimize(20, dataset, batch_size=50)):
#for i, loss in enumerate(model.LBFGS_optimize(20, dataset, lr=1, history_size=5)):
# And we liveplot the updates to the model as they happen
print(i,loss)
model.inspect(dataset)
examples/unified_modes.py
+1 -2
View File
@@ -10,8 +10,7 @@ dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
# dataset.inspect()
# plt.show()
#model = CDTools.models.UnifiedModePtycho.from_dataset(dataset, oversampling=2,n_modes=3)#, probe_support_radius=90)
model = CDTools.models.UnifiedModePtycho2.from_dataset(dataset, oversampling=1,n_modes=3)#, probe_support_radius=90)
model = CDTools.models.UnifiedModePtycho.from_dataset(dataset, oversampling=1,n_modes=3, dm_rank=-1)#, probe_support_radius=90)
#model = CDTools.models.FancyPtycho.from_dataset(dataset, oversampling=1)
model.to(device='cuda')