mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-09 21:12:42 +02:00
Merge branch 'master' of github.mit.edu:Scattering/CDTools
This commit is contained in:
Maddie Cain
@@ -7,7 +7,7 @@ from CDTools.tools import plotting as p
|
||||
from copy import copy
|
||||
from torch.utils import data as torchdata
|
||||
from matplotlib import pyplot as plt
|
||||
|
||||
import numpy as np
|
||||
|
||||
class SimplePtycho(CDIModel):
|
||||
|
||||
@@ -45,6 +45,7 @@ class SimplePtycho(CDIModel):
|
||||
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
|
||||
|
||||
|
||||
|
||||
@classmethod
|
||||
def from_dataset(cls, dataset):
|
||||
wavelength = dataset.wavelength
|
||||
@@ -138,14 +139,10 @@ class SimplePtycho(CDIModel):
|
||||
|
||||
|
||||
def inspect(self):
|
||||
p.plot_amplitude(self.probe, basis=self.probe_basis)
|
||||
plt.title('Probe Amplitude')
|
||||
p.plot_phase(self.probe, basis=self.probe_basis)
|
||||
plt.title('Probe Phase')
|
||||
p.plot_amplitude(self.obj, basis=self.probe_basis)
|
||||
plt.title('Object Amplitude')
|
||||
p.plot_phase(self.obj, basis=self.probe_basis)
|
||||
plt.title('Object Phase')
|
||||
p.plot_amplitude(self.probe, basis=self.probe_basis, title = 'Probe Amplitude')
|
||||
p.plot_phase(self.probe, basis=self.probe_basis, title = 'Probe Phase')
|
||||
p.plot_amplitude(self.obj, basis=self.probe_basis, title = 'Object Amplitude')
|
||||
p.plot_phase(self.obj, basis=self.probe_basis, title = 'Object Phase')
|
||||
|
||||
|
||||
def save_results(self):
|
||||
@@ -153,8 +150,8 @@ class SimplePtycho(CDIModel):
|
||||
probe = probe * self.probe_norm.detach().cpu().numpy()
|
||||
obj = tools.cmath.torch_to_complex(self.obj.detach().cpu())
|
||||
return {'probe':probe,'obj':obj}
|
||||
|
||||
|
||||
|
||||
|
||||
def ePIE(self, iterations, dataset, beta = 1.0):
|
||||
"""Runs an ePIE reconstruction as described in `Maiden et al. (2017) <https://www.osapublishing.org/optica/abstract.cfm?uri=optica-4-7-736>`_.
|
||||
Optional parameters are:
|
||||
@@ -172,13 +169,13 @@ class SimplePtycho(CDIModel):
|
||||
mask=None
|
||||
|
||||
def probe_update(exit_wave, exit_wave_corrected, probe, object, translation):
|
||||
return probe+tools.cmath.cmult(beta*(tools.cmath.cconj(object)/t.max(tools.cmath.cabssq(object)))[translation[0]:translation[0]+probe_shape[0],translation[1]:translation[1]+probe_shape[1]], \
|
||||
exit_wave_corrected-exit_wave)
|
||||
new_probe = probe + tools.cmath.cmult(beta * tools.cmath.cconj(object[translation])/(self.probe_norm*t.max(tools.cmath.cabssq(object))), exit_wave_corrected-exit_wave)
|
||||
return new_probe
|
||||
|
||||
def object_update(exit_wave, exit_wave_corrected, probe, object, translation):
|
||||
object[translation[0]:translation[0]+probe_shape[0],translation[1]:translation[1]+probe_shape[1]]\
|
||||
+=tools.cmath.cmult(tools.cmath.cconj(probe)/t.max(tools.cmath.cabssq(probe)),exit_wave_corrected-exit_wave)
|
||||
return object
|
||||
new_object = object.clone()
|
||||
new_object[translation] = object[translation] + tools.cmath.cmult(beta * tools.cmath.cconj(probe)/(self.probe_norm*t.max(tools.cmath.cabssq(probe))), exit_wave_corrected-exit_wave)
|
||||
return new_object
|
||||
|
||||
with t.no_grad():
|
||||
data_loader = torchdata.DataLoader(dataset, shuffle=True)
|
||||
@@ -186,15 +183,33 @@ class SimplePtycho(CDIModel):
|
||||
for it in range(iterations):
|
||||
loss = []
|
||||
for (i, [translations]), [patterns] in data_loader:
|
||||
probe = self.probe.clone()
|
||||
object = self.obj.clone()
|
||||
exit_wave = self.interaction(i, translations)
|
||||
exit_wave_corrected = exit_wave.clone()
|
||||
exit_wave_corrected[self.detector_slice] = tools.projectors.modulus(self.forward_propagator(exit_wave)[self.detector_slice], patterns, mask = mask)
|
||||
probe = self.probe.data.clone()
|
||||
object = self.obj.data.clone()
|
||||
|
||||
integer_translations = t.round(translations).to(dtype=t.int32)
|
||||
self.probe.data = probe_update(exit_wave, exit_wave_corrected, probe, object, integer_translations)
|
||||
self.obj.data = object_update(exit_wave, exit_wave_corrected, probe, object, integer_translations)
|
||||
exit_wave = self.interaction(i, translations).clone()
|
||||
# Apply modulus constraint
|
||||
exit_wave_corrected = exit_wave.clone()
|
||||
exit_wave_corrected = self.forward_propagator(exit_wave_corrected.clone())
|
||||
exit_wave_corrected[self.detector_slice] = tools.projectors.modulus(exit_wave_corrected.clone()[self.detector_slice], patterns, mask = mask)
|
||||
exit_wave_corrected = self.backward_propagator(exit_wave_corrected.clone())
|
||||
|
||||
# Calculate the section of the object wavefunction to be modified
|
||||
pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
|
||||
translations)
|
||||
pix_trans -= self.min_translation
|
||||
|
||||
pix_trans = t.round(pix_trans).to(dtype=t.int32).numpy()
|
||||
|
||||
object_slice = np.s_[pix_trans[0]:
|
||||
pix_trans[0]+probe_shape[0],
|
||||
pix_trans[1]:
|
||||
pix_trans[1]+probe_shape[1]]
|
||||
|
||||
# Apply probe and object updates
|
||||
self.probe.data = probe_update(exit_wave, exit_wave_corrected, probe, object, object_slice)
|
||||
self.obj.data = object_update(exit_wave, exit_wave_corrected, probe, object, object_slice)
|
||||
|
||||
# Calculate loss
|
||||
loss.append(self.loss(self.measurement(self.interaction(i, translations)), patterns))
|
||||
|
||||
yield t.mean(t.Tensor(loss)).cpu().numpy()
|
||||
yield it, t.mean(t.Tensor(loss)).cpu().numpy()
|
||||
|
||||
@@ -22,7 +22,7 @@ def modulus(wavefront, intensities, mask = None):
|
||||
wavefront (torch.Tensor) : The JxNxMx2 stack of complex propagated wavefronts
|
||||
intensities (torch.Tensor): The measured diffraction pattern(s) stored as an JxNxM stack of real tensors
|
||||
mask (torch.Tensor) : Mask for the intensities array with shape JxNxM, where bad detector pixels are set to 0 and usable pixels set to 1
|
||||
|
||||
|
||||
Returns:
|
||||
torch.Tensor : The JxNxMx2 propagated wavefield with corrected intensities
|
||||
"""
|
||||
@@ -55,10 +55,9 @@ def support(wavefront, support):
|
||||
|
||||
Args:
|
||||
wavefront (torch.Tensor) : The JxNxMx2 stack of complex propagated wavefronts
|
||||
support (torch.Tensor) : An NxM support, with 1s within the support and 0s outside
|
||||
|
||||
support (torch.Tensor) : An NxM support, with 1s within the support and 0s outside
|
||||
|
||||
Returns:
|
||||
torch.Tensor : The JxNxMx2 wavefield with the support mask applied
|
||||
"""
|
||||
return wavefront * support.to(wavefront.dtype)[...,None]
|
||||
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
import CDTools
|
||||
from CDTools.tools.initializers import gaussian
|
||||
from CDTools.tools.propagators import far_field
|
||||
from CDTools.tools.measurements import intensity
|
||||
|
||||
from CDTools.tools.cmath import *
|
||||
from CDTools.tools import interactions
|
||||
import h5py
|
||||
import numpy as np
|
||||
from matplotlib import pyplot as plt
|
||||
import torch as t
|
||||
import scipy.misc
|
||||
|
||||
|
||||
|
||||
def simulate_pattern(pixel_translation_step = (10, 10), probe = None, intensity_err = True,
|
||||
random_noise = True, translation_err = True, background_noise = True):
|
||||
"""
|
||||
Generates diffraction data from scipy's face image with a gaussian phase.
|
||||
Can simulate different experimental errors, including:
|
||||
* Varying probe intensities over time
|
||||
* Nanopositioner location uncertainties
|
||||
* Background detector noise
|
||||
* Random noise
|
||||
* Uncentered probe
|
||||
|
||||
We can do this by simulating the data collection process with a large
|
||||
probe array, and changing the translation by a random (integer) amount.
|
||||
Then, we can scale the intensity by a random amount (close to 1),
|
||||
replace pixels with some scaled version of their current value to simulate
|
||||
random noise, then add in horizontal and vertical detector background noise.
|
||||
Finally, we slice the simulated diffraction patterns to uncenter the probe
|
||||
on the final patterns [TO DO].
|
||||
|
||||
Args:
|
||||
pixel_translation_step (array-like) : a tuple containing (y pixel translation, x pixel translation)
|
||||
probe (np.ndarray) : a two-dimensional array representing the probe wavefunction. Defaults to a gaussian probe
|
||||
intensity_err (Boolean) : Defaults to True, to simulate varying probe intensities over time
|
||||
random_noise (Boolean) : Defaults to True, indicating adding random noise to the diffraction patterns
|
||||
translation_err (Boolean) : Defaults to True, to simulate nanopositioner errors
|
||||
background_noise (Boolean) : [TO DO]
|
||||
Returns:
|
||||
patterns (t.Tensor) : NxMx2 array of simulated diffraction patterns with error
|
||||
translations (t.Tensor) : N*Mx2 array of instrument translations
|
||||
"""
|
||||
obj = scipy.misc.face()
|
||||
# Convert to grayscale
|
||||
obj = 0.2989 * obj[:,:,0] + 0.5870 * obj[:,:,1] + 0.1140 * obj[:,:,2]
|
||||
phase = gaussian(obj.shape, [500,500], amplitude=1, center = None)
|
||||
obj = obj*np.exp(1j*torch_to_complex(phase))
|
||||
|
||||
if probe is None:
|
||||
probe = torch_to_complex(gaussian([512,512], [5,5], amplitude=1, center = None))
|
||||
|
||||
obj_shape = np.array(obj.shape)
|
||||
probe_shape = np.array(probe.shape)
|
||||
|
||||
translation_range = (obj_shape-probe_shape)//np.array(pixel_translation_step)*np.array(pixel_translation_step)
|
||||
translations = np.mgrid[0:translation_range[0]+pixel_translation_step[0]:pixel_translation_step[0], 0:translation_range[1]+pixel_translation_step[1]:pixel_translation_step[1]]
|
||||
translations = translations.T.reshape((translations.T.shape[0]*translations.T.shape[1], 2))
|
||||
translations = t.tensor(translations, dtype = t.float32)
|
||||
ideal_translations = t.tensor(translations, dtype = t.float32)
|
||||
obj = complex_to_torch(obj)
|
||||
probe = complex_to_torch(probe)
|
||||
|
||||
if translation_err:
|
||||
err = t.normal(mean=t.zeros(translations.shape), std=t.ones(translations.shape)/2)
|
||||
translations += err
|
||||
# Ensure that the new translations don't go off the object area
|
||||
translations[translations<0] = 0
|
||||
translation_range = t.tensor(translation_range, dtype = t.float32)
|
||||
translations[...,0][translations[...,0] > translation_range[0]] = translation_range[0]
|
||||
translations[...,1][translations[...,1] > translation_range[1]] = translation_range[1]
|
||||
|
||||
patterns = intensity(far_field(interactions.ptycho_2D_round(probe, obj, translations)))
|
||||
|
||||
if intensity_err:
|
||||
patterns *= t.normal(mean = t.ones(patterns.shape), std = t.ones(patterns.shape)/32)
|
||||
|
||||
if background_noise:
|
||||
patterns += t.normal(mean = t.zeeros(patterns.shape), std = t.ones(patterns.shape)*t.max(patterns)/32)
|
||||
|
||||
|
||||
return patterns, ideal_translations
|
||||
@@ -12,8 +12,8 @@ import pickle
|
||||
|
||||
#filename = '../../Downloads/114429_p.cxi'
|
||||
#filename = '../../../Projects/CSX_3_19/cxis/processed/114429_p.cxi'
|
||||
#filename = '../../../Projects/CSX_3_19/cxis/processed/115145_p.cxi'
|
||||
filename = '../../../Downloads/AuBalls_700ms_30nmStep_3_3SS_filter.cxi'
|
||||
filename = '../../../Projects/CSX_3_19/cxis/processed/115145_p.cxi'
|
||||
#filename = '../../Desktop/Reconstructions/114429_p.cxi'
|
||||
|
||||
|
||||
with h5py.File(filename,'r') as f:
|
||||
@@ -24,16 +24,19 @@ model = CDTools.models.SimplePtycho.from_dataset(dataset)
|
||||
|
||||
# Uncomment these to use on the CPU
|
||||
# default is CPU with 32-bit floats
|
||||
model.to(device='cuda')
|
||||
model.to(device='cpu')
|
||||
#dataset.to(device='cuda')
|
||||
dataset.get_as(device='cuda')
|
||||
dataset.get_as(device='cpu')
|
||||
|
||||
for i, loss in enumerate(model.ePIE(10, dataset)):
|
||||
print(i, loss)
|
||||
|
||||
for loss in model.Adam_optimize(1, dataset):
|
||||
print(loss)
|
||||
for i, loss in enumerate(model.Adam_optimize(10, dataset)):
|
||||
print(i, loss)
|
||||
|
||||
model.inspect()
|
||||
|
||||
#with open('test_results.pickle', 'wb') as f:
|
||||
# pickle.dump(model.save_results(),f)
|
||||
|
||||
model.inspect()
|
||||
plt.show()
|
||||
|
||||
Reference in New Issue
Block a user